




import inspect
import unittest

from datasets import load_dataset
from packaging import version

from transformers import BeitConfig
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import (
           MODEL_MAPPING,
           BeitForImageClassification,
           BeitForMaskedImageModeling,
           BeitForSemanticSegmentation,
           BeitModel,
       )
       from transformers.models.beit.modeling_beit import BEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       import PIL
       from PIL import Image

       from transformers import BeitImageProcessor




class 		UpperCamelCase__	:

    '''simple docstring'''



    def __init__(    self     ,  UpperCamelCase__     ,  UpperCamelCase__=100     ,  UpperCamelCase__=13     ,  UpperCamelCase__=30     ,  UpperCamelCase__=2     ,  UpperCamelCase__=3     ,  UpperCamelCase__=True     ,  UpperCamelCase__=True     ,  UpperCamelCase__=32     ,  UpperCamelCase__=4     ,  UpperCamelCase__=4     ,  UpperCamelCase__=37     ,  UpperCamelCase__="gelu"     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=0.1     ,  UpperCamelCase__=10     ,  UpperCamelCase__=0.02     ,  UpperCamelCase__=3     ,  UpperCamelCase__=None     ,  UpperCamelCase__=[0, 1, 2, 3]     ,  )     ->      List[Any]:
           lowerCamelCase   :      Optional[Any]			 =					parent
           lowerCamelCase   :      str			 =					100
           lowerCamelCase   :      int			 =					batch_size
           lowerCamelCase   :      Tuple			 =					image_size
           lowerCamelCase   :      List[str]			 =					patch_size
           lowerCamelCase   :      Dict			 =					num_channels
           lowerCamelCase   :      Optional[int]			 =					is_training
           lowerCamelCase   :      int			 =					use_labels
           lowerCamelCase   :      Optional[int]			 =					hidden_size
           lowerCamelCase   :      Optional[int]			 =					num_hidden_layers
           lowerCamelCase   :      List[Any]			 =					num_attention_heads
           lowerCamelCase   :      Any			 =					intermediate_size
           lowerCamelCase   :      List[Any]			 =					hidden_act
           lowerCamelCase   :      str			 =					hidden_dropout_prob
           lowerCamelCase   :      Tuple			 =					attention_probs_dropout_prob
           lowerCamelCase   :      Dict			 =					type_sequence_label_size
           lowerCamelCase   :      int			 =					initializer_range
           lowerCamelCase   :      Optional[int]			 =					scope
           lowerCamelCase   :      Tuple			 =					out_indices
           lowerCamelCase   :      Dict			 =					num_labels

           # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
           lowerCamelCase   :      List[str]			 =					(image_size // patch_size) ** 2
           lowerCamelCase   :      Optional[Any]			 =					num_patches + 1







    def     _lowercase     (    self     )     ->      List[Any]:
           lowerCamelCase   :      Any			 =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]     )

           lowerCamelCase   :      Optional[int]			 =					None
           lowerCamelCase   :      Dict			 =					None
           if self.use_labels:
                  lowerCamelCase   :      Tuple			 =					ids_tensor([self.batch_size]     ,  self.type_sequence_label_size     )
                  lowerCamelCase   :      Tuple			 =					ids_tensor([self.batch_size, self.image_size, self.image_size]     ,  self.num_labels     )

           lowerCamelCase   :      List[str]			 =					self.get_config()

           return config, pixel_values, labels, pixel_labels







    def     _lowercase     (    self     )     ->      str:
           return BeitConfig(
               vocab_size=self.vocab_size     ,  image_size=self.image_size     ,  patch_size=self.patch_size     ,  num_channels=self.num_channels     ,  hidden_size=self.hidden_size     ,  num_hidden_layers=self.num_hidden_layers     ,  num_attention_heads=self.num_attention_heads     ,  intermediate_size=self.intermediate_size     ,  hidden_act=self.hidden_act     ,  hidden_dropout_prob=self.hidden_dropout_prob     ,  attention_probs_dropout_prob=self.attention_probs_dropout_prob     ,  is_decoder=UpperCamelCase__     ,  initializer_range=self.initializer_range     ,  out_indices=self.out_indices     ,  )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      Any:
           lowerCamelCase   :      Dict			 =					BeitModel(config=UpperCamelCase__     )
           model.to(UpperCamelCase__     )
           model.eval()
           lowerCamelCase   :      Optional[Any]			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(result.last_hidden_state.shape     ,  (self.batch_size, self.seq_length, self.hidden_size)     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      Optional[int]:
           lowerCamelCase   :      int			 =					BeitForMaskedImageModeling(config=UpperCamelCase__     )
           model.to(UpperCamelCase__     )
           model.eval()
           lowerCamelCase   :      Tuple			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.seq_length - 1, self.vocab_size)     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      Tuple:
           lowerCamelCase   :      str			 =					self.type_sequence_label_size
           lowerCamelCase   :      Optional[Any]			 =					BeitForImageClassification(UpperCamelCase__     )
           model.to(UpperCamelCase__     )
           model.eval()
           lowerCamelCase   :      List[Any]			 =					model(UpperCamelCase__     ,  labels=UpperCamelCase__     )
           self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.type_sequence_label_size)     )

           # test greyscale images
           lowerCamelCase   :      List[Any]			 =					1
           lowerCamelCase   :      Any			 =					BeitForImageClassification(UpperCamelCase__     )
           model.to(UpperCamelCase__     )
           model.eval()

           lowerCamelCase   :      List[str]			 =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]     )
           lowerCamelCase   :      Any			 =					model(UpperCamelCase__     ,  labels=UpperCamelCase__     )
           self.parent.assertEqual(result.logits.shape     ,  (self.batch_size, self.type_sequence_label_size)     )







    def     _lowercase     (    self     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     ,  UpperCamelCase__     )     ->      List[str]:
           lowerCamelCase   :      Optional[int]			 =					self.num_labels
           lowerCamelCase   :      str			 =					BeitForSemanticSegmentation(UpperCamelCase__     )
           model.to(UpperCamelCase__     )
           model.eval()
           lowerCamelCase   :      Tuple			 =					model(UpperCamelCase__     )
           self.parent.assertEqual(
               result.logits.shape     ,  (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )
           lowerCamelCase   :      Optional[int]			 =					model(UpperCamelCase__     ,  labels=UpperCamelCase__     )
           self.parent.assertEqual(
               result.logits.shape     ,  (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)     )







    def     _lowercase     (    self     )     ->      Optional[Any]:
           lowerCamelCase   :      int			 =					self.prepare_config_and_inputs()
           lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase ,      lowerCamelCase   :      List[str]			 =					config_and_inputs
           lowerCamelCase   :      List[Any]			 =					{"pixel_values": pixel_values}
           return config, inputs_dict




@require_torch
class 		UpperCamelCase__	(lowerCAmelCase__    ,				lowerCAmelCase__    ,				unittest.TestCase     ):

    '''simple docstring'''



    lowerCamelCase_      :	Optional[int]	     =      (
        (BeitModel, BeitForImageClassification, BeitForMaskedImageModeling, BeitForSemanticSegmentation)
        if is_torch_available()
        else ()
    )
    lowerCamelCase_      :	int	     =      (
        {
            """feature-extraction""": BeitModel,
            """image-classification""": BeitForImageClassification,
            """image-segmentation""": BeitForSemanticSegmentation,
        }
        if is_torch_available()
        else {}
    )

    lowerCamelCase_      :	Tuple	     =      False
    lowerCamelCase_      :	List[str]	     =      False
    lowerCamelCase_      :	str	     =      False







    def     _lowercase     (    self     )     ->      List[str]:
           lowerCamelCase   :      Dict			 =					BeitModelTester(self     )
           lowerCamelCase   :      List[Any]			 =					ConfigTester(self     ,  config_class=UpperCamelCase__     ,  has_text_modality=UpperCamelCase__     ,  hidden_size=37     )







    def     _lowercase     (    self     )     ->      Dict:
           self.config_tester.run_common_tests()







    @unittest.skip(reason="BEiT does not use inputs_embeds"     )
    def     _lowercase     (    self     )     ->      str:
           pass







    @require_torch_multi_gpu
    @unittest.skip(reason="BEiT has some layers using `add_module` which doesn't work well with `nn.DataParallel`"     )
    def     _lowercase     (    self     )     ->      Any:
           pass







    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase ,      lowerCamelCase   :      List[str]			 =					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  lowerCamelCase   :      List[str]			 =					model_class(UpperCamelCase__     )
                  self.assertIsInstance(model.get_input_embeddings()     ,  (nn.Module)     )
                  lowerCamelCase   :      List[Any]			 =					model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(UpperCamelCase__     ,  nn.Linear     )     )







    def     _lowercase     (    self     )     ->      Tuple:
           lowerCamelCase ,      lowerCamelCase   :      Optional[int]			 =					self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  lowerCamelCase   :      Optional[Any]			 =					model_class(UpperCamelCase__     )
                  lowerCamelCase   :      Union[str, Any]			 =					inspect.signature(model.forward     )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  lowerCamelCase   :      Optional[Any]			 =					[*signature.parameters.keys()]

                  lowerCamelCase   :      str			 =					["pixel_values"]
                  self.assertListEqual(arg_names[:1]     ,  UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      Dict			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      Any:
           lowerCamelCase   :      Optional[int]			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      List[str]:
           lowerCamelCase   :      int			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      Optional[Any]:
           lowerCamelCase   :      Optional[Any]			 =					self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*UpperCamelCase__     )







    def     _lowercase     (    self     )     ->      Tuple:
           if not self.model_tester.is_training:
                  return

           lowerCamelCase ,      lowerCamelCase   :      Optional[Any]			 =					self.model_tester.prepare_config_and_inputs_for_common()
           lowerCamelCase   :      Tuple			 =					True

           for model_class in self.all_model_classes:
                  # we don't test BeitForMaskedImageModeling
                  if model_class in [*get_values(UpperCamelCase__     ), BeitForMaskedImageModeling]:
                         continue

                  lowerCamelCase   :      int			 =					model_class(UpperCamelCase__     )
                  model.to(UpperCamelCase__     )
                  model.train()
                  lowerCamelCase   :      Dict			 =					self._prepare_for_class(UpperCamelCase__     ,  UpperCamelCase__     ,  return_labels=UpperCamelCase__     )
                  lowerCamelCase   :      Any			 =					model(**UpperCamelCase__     ).loss
                  loss.backward()







    def     _lowercase     (    self     )     ->      List[str]:
           lowerCamelCase ,      lowerCamelCase   :      List[Any]			 =					self.model_tester.prepare_config_and_inputs_for_common()
           if not self.model_tester.is_training:
                  return

           lowerCamelCase   :      Optional[Any]			 =					False
           lowerCamelCase   :      str			 =					True

           for model_class in self.all_model_classes:
                  # we don't test BeitForMaskedImageModeling
                  if (
                      model_class in [*get_values(UpperCamelCase__     ), BeitForMaskedImageModeling]
                      or not model_class.supports_gradient_checkpointing
                  ):
                         continue

                  lowerCamelCase   :      Dict			 =					model_class(UpperCamelCase__     )
                  model.gradient_checkpointing_enable()
                  model.to(UpperCamelCase__     )
                  model.train()
                  lowerCamelCase   :      Union[str, Any]			 =					self._prepare_for_class(UpperCamelCase__     ,  UpperCamelCase__     ,  return_labels=UpperCamelCase__     )
                  lowerCamelCase   :      Optional[int]			 =					model(**UpperCamelCase__     ).loss
                  loss.backward()







    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase ,      lowerCamelCase   :      Any			 =					self.model_tester.prepare_config_and_inputs_for_common()

           lowerCamelCase   :      str			 =					_config_zero_init(UpperCamelCase__     )
           for model_class in self.all_model_classes:
                  lowerCamelCase   :      Any			 =					model_class(config=UpperCamelCase__     )
                  for name, param in model.named_parameters():
                         # we skip lambda parameters as these require special initial values
                         # determined by config.layer_scale_init_value
                         if "lambda" in name:
                                continue
                         if param.requires_grad:
                                self.assertIn(
                                    ((param.data.mean() * 1e9).round() / 1e9).item()     ,  [0.0, 1.0]     ,  msg=F'''Parameter {name} of model {model_class} seems not properly initialized'''     ,  )







    @slow
    def     _lowercase     (    self     )     ->      List[str]:
           for model_name in BEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  lowerCamelCase   :      Optional[Any]			 =					BeitModel.from_pretrained(UpperCamelCase__     )
                  self.assertIsNotNone(UpperCamelCase__     )




def      A      (      )	-> Tuple:
       lowerCamelCase   :      int			 =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"    )
       return image




@require_torch
@require_vision
class 		UpperCamelCase__	(unittest.TestCase     ):

    '''simple docstring'''



    @cached_property
    def     _lowercase     (    self     )     ->      Tuple:
           return BeitImageProcessor.from_pretrained("microsoft/beit-base-patch16-224"     ) if is_vision_available() else None







    @slow
    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      List[str]			 =					BeitForMaskedImageModeling.from_pretrained("microsoft/beit-base-patch16-224-pt22k"     ).to(UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					self.default_image_processor
           lowerCamelCase   :      Union[str, Any]			 =					prepare_img()
           lowerCamelCase   :      str			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="pt"     ).pixel_values.to(UpperCamelCase__     )

           # prepare bool_masked_pos
           lowerCamelCase   :      List[Any]			 =					torch.ones((1, 196)     ,  dtype=torch.bool     ).to(UpperCamelCase__     )

           # forward pass
           with torch.no_grad():
                  lowerCamelCase   :      Dict			 =					model(pixel_values=UpperCamelCase__     ,  bool_masked_pos=UpperCamelCase__     )
           lowerCamelCase   :      List[str]			 =					outputs.logits

           # verify the logits
           lowerCamelCase   :      Dict			 =					torch.Size((1, 196, 8192)     )
           self.assertEqual(logits.shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      int			 =					torch.tensor(
               [[-3.2437, 0.5072, -13.9174], [-3.2456, 0.4948, -13.9401], [-3.2033, 0.5121, -13.8550]]     ).to(UpperCamelCase__     )

           self.assertTrue(torch.allclose(logits[bool_masked_pos][:3, :3]     ,  UpperCamelCase__     ,  atol=1e-2     )     )







    @slow
    def     _lowercase     (    self     )     ->      Dict:
           lowerCamelCase   :      int			 =					BeitForImageClassification.from_pretrained("microsoft/beit-base-patch16-224"     ).to(UpperCamelCase__     )

           lowerCamelCase   :      Dict			 =					self.default_image_processor
           lowerCamelCase   :      Optional[Any]			 =					prepare_img()
           lowerCamelCase   :      int			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="pt"     ).to(UpperCamelCase__     )

           # forward pass
           with torch.no_grad():
                  lowerCamelCase   :      List[str]			 =					model(**UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					outputs.logits

           # verify the logits
           lowerCamelCase   :      Optional[int]			 =					torch.Size((1, 1000)     )
           self.assertEqual(logits.shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      str			 =					torch.tensor([-1.2385, -1.0987, -1.0108]     ).to(UpperCamelCase__     )

           self.assertTrue(torch.allclose(logits[0, :3]     ,  UpperCamelCase__     ,  atol=1e-4     )     )

           lowerCamelCase   :      List[str]			 =					281
           self.assertEqual(logits.argmax(-1     ).item()     ,  UpperCamelCase__     )







    @slow
    def     _lowercase     (    self     )     ->      str:
           lowerCamelCase   :      Union[str, Any]			 =					BeitForImageClassification.from_pretrained("microsoft/beit-large-patch16-224-pt22k-ft22k"     ).to(
               UpperCamelCase__     )

           lowerCamelCase   :      int			 =					self.default_image_processor
           lowerCamelCase   :      Tuple			 =					prepare_img()
           lowerCamelCase   :      List[Any]			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="pt"     ).to(UpperCamelCase__     )

           # forward pass
           with torch.no_grad():
                  lowerCamelCase   :      Any			 =					model(**UpperCamelCase__     )
           lowerCamelCase   :      Tuple			 =					outputs.logits

           # verify the logits
           lowerCamelCase   :      int			 =					torch.Size((1, 2_1841)     )
           self.assertEqual(logits.shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					torch.tensor([1.6881, -0.2787, 0.5901]     ).to(UpperCamelCase__     )

           self.assertTrue(torch.allclose(logits[0, :3]     ,  UpperCamelCase__     ,  atol=1e-4     )     )

           lowerCamelCase   :      str			 =					2396
           self.assertEqual(logits.argmax(-1     ).item()     ,  UpperCamelCase__     )







    @slow
    def     _lowercase     (    self     )     ->      int:
           lowerCamelCase   :      int			 =					BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"     )
           lowerCamelCase   :      str			 =					model.to(UpperCamelCase__     )

           lowerCamelCase   :      Tuple			 =					BeitImageProcessor(do_resize=UpperCamelCase__     ,  size=640     ,  do_center_crop=UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					load_dataset("hf-internal-testing/fixtures_ade20k"     ,  split="test"     )
           lowerCamelCase   :      str			 =					Image.open(ds[0]["file"]     )
           lowerCamelCase   :      List[Any]			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="pt"     ).to(UpperCamelCase__     )

           # forward pass
           with torch.no_grad():
                  lowerCamelCase   :      Optional[int]			 =					model(**UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					outputs.logits

           # verify the logits
           lowerCamelCase   :      List[str]			 =					torch.Size((1, 150, 160, 160)     )
           self.assertEqual(logits.shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      Any			 =					version.parse(PIL.__version__     ) < version.parse("9.0.0"     )

           if is_pillow_less_than_a:
                  lowerCamelCase   :      List[str]			 =					torch.tensor(
                      [
                          [[-4.9225, -2.3954, -3.0522], [-2.8822, -1.0046, -1.7561], [-2.9549, -1.3228, -2.1347]],
                          [[-5.8168, -3.4129, -4.0778], [-3.8651, -2.2214, -3.0277], [-3.8356, -2.4643, -3.3535]],
                          [[-0.0078, 3.9952, 4.0754], [2.9856, 4.6944, 5.0035], [3.2413, 4.7813, 4.9969]],
                      ]     ,  device=UpperCamelCase__     ,  )
           else:
                  lowerCamelCase   :      str			 =					torch.tensor(
                      [
                          [[-4.8960, -2.3688, -3.0355], [-2.8478, -0.9836, -1.7418], [-2.9449, -1.3332, -2.1456]],
                          [[-5.8081, -3.4124, -4.1006], [-3.8561, -2.2081, -3.0323], [-3.8365, -2.4601, -3.3669]],
                          [[-0.0309, 3.9868, 4.0540], [2.9640, 4.6877, 4.9976], [3.2081, 4.7690, 4.9942]],
                      ]     ,  device=UpperCamelCase__     ,  )

           self.assertTrue(torch.allclose(logits[0, :3, :3, :3]     ,  UpperCamelCase__     ,  atol=1e-4     )     )







    @slow
    def     _lowercase     (    self     )     ->      List[str]:
           lowerCamelCase   :      List[Any]			 =					BeitForSemanticSegmentation.from_pretrained("microsoft/beit-base-finetuned-ade-640-640"     )
           lowerCamelCase   :      Optional[int]			 =					model.to(UpperCamelCase__     )

           lowerCamelCase   :      int			 =					BeitImageProcessor(do_resize=UpperCamelCase__     ,  size=640     ,  do_center_crop=UpperCamelCase__     )

           lowerCamelCase   :      List[Any]			 =					load_dataset("hf-internal-testing/fixtures_ade20k"     ,  split="test"     )
           lowerCamelCase   :      str			 =					Image.open(ds[0]["file"]     )
           lowerCamelCase   :      Any			 =					image_processor(images=UpperCamelCase__     ,  return_tensors="pt"     ).to(UpperCamelCase__     )

           # forward pass
           with torch.no_grad():
                  lowerCamelCase   :      Dict			 =					model(**UpperCamelCase__     )

           lowerCamelCase   :      Optional[int]			 =					outputs.logits.detach().cpu()

           lowerCamelCase   :      int			 =					image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase__     ,  target_sizes=[(500, 300)]     )
           lowerCamelCase   :      int			 =					torch.Size((500, 300)     )
           self.assertEqual(segmentation[0].shape     ,  UpperCamelCase__     )

           lowerCamelCase   :      List[str]			 =					image_processor.post_process_semantic_segmentation(outputs=UpperCamelCase__     )
           lowerCamelCase   :      Optional[Any]			 =					torch.Size((160, 160)     )
           self.assertEqual(segmentation[0].shape     ,  UpperCamelCase__     )








# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


UpperCAmelCase	      = '''Create a default config file for Accelerate with only a few flags set.'''






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE="no" ,							__SCREAMING_SNAKE_CASE = default_json_config_file ,							__SCREAMING_SNAKE_CASE = False							):
	lowercase       =    Path(__SCREAMING_SNAKE_CASE							)
	path.parent.mkdir(parents=__SCREAMING_SNAKE_CASE ,							exist_ok=__SCREAMING_SNAKE_CASE							)
	if path.exists():
		print(
		    F'''Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'''							)
		return False
	lowercase       =    mixed_precision.lower()
	if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
		raise ValueError(
		    F'''`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'''							)
	lowercase       =    {
	    'compute_environment': 'LOCAL_MACHINE',
	    'mixed_precision': mixed_precision,
	}
	if torch.cuda.is_available():
		lowercase       =    torch.cuda.device_count()
		lowercase       =    num_gpus
		lowercase       =    False
		if num_gpus > 1:
			lowercase       =    'MULTI_GPU'
		else:
			lowercase       =    'NO'
	elif is_xpu_available() and use_xpu:
		lowercase       =    torch.xpu.device_count()
		lowercase       =    num_xpus
		lowercase       =    False
		if num_xpus > 1:
			lowercase       =    'MULTI_XPU'
		else:
			lowercase       =    'NO'
	elif is_npu_available():
		lowercase       =    torch.npu.device_count()
		lowercase       =    num_npus
		lowercase       =    False
		if num_npus > 1:
			lowercase       =    'MULTI_NPU'
		else:
			lowercase       =    'NO'
	else:
		lowercase       =    0
		lowercase       =    True
		lowercase       =    1
		lowercase       =    'NO'
	lowercase       =    ClusterConfig(**__SCREAMING_SNAKE_CASE							)
	config.to_json_file(__SCREAMING_SNAKE_CASE							)
	return path






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE ,							__SCREAMING_SNAKE_CASE							):
	lowercase       =    parser.add_parser('default' ,							parents=__SCREAMING_SNAKE_CASE ,							help=__SCREAMING_SNAKE_CASE ,							formatter_class=__SCREAMING_SNAKE_CASE							)
	parser.add_argument(
	    '--config_file' ,							default=__SCREAMING_SNAKE_CASE ,							help=(
	        'The path to use to store the config file. Will default to a file named default_config.yaml in the cache '
	        'location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '
	        'such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '
	        'with \'huggingface\'.'
	    ) ,							dest='save_location' ,							)

	parser.add_argument(
	    '--mixed_precision' ,							choices=['no', 'fp16', 'bf16'] ,							type=__SCREAMING_SNAKE_CASE ,							help='Whether or not to use mixed precision training. '
	    'Choose between FP16 and BF16 (bfloat16) training. '
	    'BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.' ,							default='no' ,							)
	parser.set_defaults(func=__SCREAMING_SNAKE_CASE							)
	return parser






def 							UpperCAmelCase_		(							__SCREAMING_SNAKE_CASE							):
	lowercase       =    write_basic_config(args.mixed_precision ,							args.save_location							)
	if config_file:
		print(F'''accelerate configuration saved at {config_file}'''							)

