





'''simple docstring'''






import inspect
import unittest
import warnings

from transformers import DeiTConfig
from transformers.models.auto import get_values
from transformers.testing_utils import (
    require_accelerate,
    require_torch,
    require_torch_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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch
      from torch import nn

      from transformers import (
          MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING,
          MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
          MODEL_MAPPING,
          DeiTForImageClassification,
          DeiTForImageClassificationWithTeacher,
          DeiTForMaskedImageModeling,
          DeiTModel,
      )
      from transformers.models.deit.modeling_deit import DEIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
      from PIL import Image

      from transformers import DeiTImageProcessor




class 			UpperCAmelCase					:




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

            # in DeiT, the seq length equals the number of patches + 2 (we add 2 for the [CLS] and distilation tokens)
            _lowerCAmelCase												=							(image_size // patch_size) ** 2
            _lowerCAmelCase												=							num_patches + 2




      def     lowercase__   (      self      :			str  )    ->				int:
            _lowerCAmelCase												=							floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

            _lowerCAmelCase												=							None
            if self.use_labels:
                  _lowerCAmelCase												=							ids_tensor([self.batch_size]      ,      self.type_sequence_label_size  )

            _lowerCAmelCase												=							self.get_config()

            return config, pixel_values, labels




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Optional[int]:
            return DeiTConfig(
                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=a__      ,      initializer_range=self.initializer_range      ,      encoder_stride=self.encoder_stride      ,      )




      def     lowercase__   (      self      :			Tuple      ,      __snake_case      :			Any      ,      __snake_case      :			Union[str, Any]      ,      __snake_case      :			Optional[Any]  )    ->				Optional[int]:
            _lowerCAmelCase												=							DeiTModel(config=a__  )
            model.to(a__  )
            model.eval()
            _lowerCAmelCase												=							model(a__  )
            self.parent.assertEqual(result.last_hidden_state.shape      ,      (self.batch_size, self.seq_length, self.hidden_size)  )




      def     lowercase__   (      self      :			str      ,      __snake_case      :			Optional[int]      ,      __snake_case      :			int      ,      __snake_case      :			Optional[Any]  )    ->				int:
            _lowerCAmelCase												=							DeiTForMaskedImageModeling(config=a__  )
            model.to(a__  )
            model.eval()
            _lowerCAmelCase												=							model(a__  )
            self.parent.assertEqual(
                result.reconstruction.shape      ,      (self.batch_size, self.num_channels, self.image_size, self.image_size)  )

            # test greyscale images
            _lowerCAmelCase												=							1
            _lowerCAmelCase												=							DeiTForMaskedImageModeling(a__  )
            model.to(a__  )
            model.eval()

            _lowerCAmelCase												=							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
            _lowerCAmelCase												=							model(a__  )
            self.parent.assertEqual(result.reconstruction.shape      ,      (self.batch_size, 1, self.image_size, self.image_size)  )




      def     lowercase__   (      self      :			Optional[int]      ,      __snake_case      :			List[str]      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Optional[int]  )    ->				List[Any]:
            _lowerCAmelCase												=							self.type_sequence_label_size
            _lowerCAmelCase												=							DeiTForImageClassification(a__  )
            model.to(a__  )
            model.eval()
            _lowerCAmelCase												=							model(a__      ,      labels=a__  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.type_sequence_label_size)  )

            # test greyscale images
            _lowerCAmelCase												=							1
            _lowerCAmelCase												=							DeiTForImageClassification(a__  )
            model.to(a__  )
            model.eval()

            _lowerCAmelCase												=							floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
            _lowerCAmelCase												=							model(a__      ,      labels=a__  )
            self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.type_sequence_label_size)  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Tuple:
            _lowerCAmelCase												=							self.prepare_config_and_inputs()
            (
                (
                _lowerCAmelCase				
            )					,							(
                _lowerCAmelCase				
            )					,							(
                _lowerCAmelCase				
            )					,											
            )												=							config_and_inputs
            _lowerCAmelCase												=							{"""pixel_values""": pixel_values}
            return config, inputs_dict




@require_torch
class 			UpperCAmelCase					(      lowercase_							,		lowercase_							,		unittest.TestCase       ):
      _lowercase:      List[Any]											=       (
          (
              DeiTModel,
              DeiTForImageClassification,
              DeiTForImageClassificationWithTeacher,
              DeiTForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      _lowercase:      List[str]											=       (
          {
              "feature-extraction": DeiTModel,
              "image-classification": (DeiTForImageClassification, DeiTForImageClassificationWithTeacher),
          }
          if is_torch_available()
          else {}
      )

      _lowercase:      Union[str, Any]											=       False
      _lowercase:      List[Any]											=       False
      _lowercase:      int											=       False




      def     lowercase__   (      self      :			int  )    ->				Any:
            _lowerCAmelCase												=							DeiTModelTester(self  )
            _lowerCAmelCase												=							ConfigTester(self      ,      config_class=a__      ,      has_text_modality=a__      ,      hidden_size=37  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Union[str, Any]:
            self.config_tester.run_common_tests()




      @unittest.skip(reason="""DeiT does not use inputs_embeds"""  )
      def     lowercase__   (      self      :			List[Any]  )    ->				str:
            pass




      def     lowercase__   (      self      :			List[str]  )    ->				Union[str, Any]:
            _lowerCAmelCase					,							_lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  _lowerCAmelCase												=							model_class(a__  )
                  self.assertIsInstance(model.get_input_embeddings()      ,      (nn.Module)  )
                  _lowerCAmelCase												=							model.get_output_embeddings()
                  self.assertTrue(x is None or isinstance(a__      ,      nn.Linear  )  )




      def     lowercase__   (      self      :			Optional[int]  )    ->				Any:
            _lowerCAmelCase					,							_lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_common()

            for model_class in self.all_model_classes:
                  _lowerCAmelCase												=							model_class(a__  )
                  _lowerCAmelCase												=							inspect.signature(model.forward  )
                  # signature.parameters is an OrderedDict => so arg_names order is deterministic
                  _lowerCAmelCase												=							[*signature.parameters.keys()]

                  _lowerCAmelCase												=							["""pixel_values"""]
                  self.assertListEqual(arg_names[:1]      ,      a__  )




      def     lowercase__   (      self      :			str  )    ->				int:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_model(*a__  )




      def     lowercase__   (      self      :			str  )    ->				List[str]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_masked_image_modeling(*a__  )




      def     lowercase__   (      self      :			Union[str, Any]  )    ->				List[Any]:
            _lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_for_image_classification(*a__  )




      def     lowercase__   (      self      :			Dict      ,      __snake_case      :			List[Any]      ,      __snake_case      :			List[Any]      ,      __snake_case      :			Optional[int]=False  )    ->				str:
            _lowerCAmelCase												=							super()._prepare_for_class(a__      ,      a__      ,      return_labels=a__  )

            if return_labels:
                  if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
                        del inputs_dict["labels"]

            return inputs_dict




      def     lowercase__   (      self      :			Any  )    ->				Dict:




            if not self.model_tester.is_training:
                  return

            _lowerCAmelCase					,							_lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_common()
            _lowerCAmelCase												=							True

            for model_class in self.all_model_classes:
                  # DeiTForImageClassificationWithTeacher supports inference-only
                  if (
                      model_class in get_values(a__  )
                      or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                  ):
                        continue
                  _lowerCAmelCase												=							model_class(a__  )
                  model.to(a__  )
                  model.train()
                  _lowerCAmelCase												=							self._prepare_for_class(a__      ,      a__      ,      return_labels=a__  )
                  _lowerCAmelCase												=							model(**a__  ).loss
                  loss.backward()




      def     lowercase__   (      self      :			Dict  )    ->				Tuple:
            _lowerCAmelCase					,							_lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_common()
            if not self.model_tester.is_training:
                  return

            _lowerCAmelCase												=							False
            _lowerCAmelCase												=							True

            for model_class in self.all_model_classes:
                  if model_class in get_values(a__  ) or not model_class.supports_gradient_checkpointing:
                        continue
                  # DeiTForImageClassificationWithTeacher supports inference-only
                  if model_class.__name__ == "DeiTForImageClassificationWithTeacher":
                        continue
                  _lowerCAmelCase												=							model_class(a__  )
                  model.gradient_checkpointing_enable()
                  model.to(a__  )
                  model.train()
                  _lowerCAmelCase												=							self._prepare_for_class(a__      ,      a__      ,      return_labels=a__  )
                  _lowerCAmelCase												=							model(**a__  ).loss
                  loss.backward()




      def     lowercase__   (      self      :			Optional[int]  )    ->				str:
            _lowerCAmelCase					,							_lowerCAmelCase												=							self.model_tester.prepare_config_and_inputs_for_common()

            _lowerCAmelCase												=							[
                {"""title""": """multi_label_classification""", """num_labels""": 2, """dtype""": torch.float},
                {"""title""": """single_label_classification""", """num_labels""": 1, """dtype""": torch.long},
                {"""title""": """regression""", """num_labels""": 1, """dtype""": torch.float},
            ]

            for model_class in self.all_model_classes:
                  if (
                      model_class
                      not in [
                          *get_values(a__  ),
                          *get_values(a__  ),
                      ]
                      or model_class.__name__ == "DeiTForImageClassificationWithTeacher"
                  ):
                        continue

                  for problem_type in problem_types:
                        with self.subTest(msg=f"Testing {model_class} with {problem_type['title']}"  ):
                              _lowerCAmelCase												=							problem_type["""title"""]
                              _lowerCAmelCase												=							problem_type["""num_labels"""]

                              _lowerCAmelCase												=							model_class(a__  )
                              model.to(a__  )
                              model.train()

                              _lowerCAmelCase												=							self._prepare_for_class(a__      ,      a__      ,      return_labels=a__  )

                              if problem_type["num_labels"] > 1:
                                    _lowerCAmelCase												=							inputs["""labels"""].unsqueeze(1  ).repeat(1      ,      problem_type["""num_labels"""]  )

                              _lowerCAmelCase												=							inputs["""labels"""].to(problem_type["""dtype"""]  )

                              # This tests that we do not trigger the warning form PyTorch "Using a target size that is different
                              # to the input size. This will likely lead to incorrect results due to broadcasting. Please ensure
                              # they have the same size." which is a symptom something in wrong for the regression problem.
                              # See https://github.com/huggingface/transformers/issues/11780
                              with warnings.catch_warnings(record=a__  ) as warning_list:
                                    _lowerCAmelCase												=							model(**a__  ).loss
                              for w in warning_list:
                                    if "Using a target size that is different to the input size" in str(w.message  ):
                                          raise ValueError(
                                              f"Something is going wrong in the regression problem: intercepted {w.message}"  )

                              loss.backward()




      @slow
      def     lowercase__   (      self      :			List[str]  )    ->				Optional[int]:




            for model_name in DEIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  _lowerCAmelCase												=							DeiTModel.from_pretrained(a__  )
                  self.assertIsNotNone(a__  )




def        UpperCamelCase__				(		):





      """simple docstring"""


      _lowerCAmelCase												=							Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""  )
      return image





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




      @cached_property
      def     lowercase__   (      self      :			Optional[Any]  )    ->				Optional[int]:
            return (
                DeiTImageProcessor.from_pretrained("""facebook/deit-base-distilled-patch16-224"""  )
                if is_vision_available()
                else None
            )




      @slow
      def     lowercase__   (      self      :			Any  )    ->				Dict:
            _lowerCAmelCase												=							DeiTForImageClassificationWithTeacher.from_pretrained("""facebook/deit-base-distilled-patch16-224"""  ).to(
                a__  )

            _lowerCAmelCase												=							self.default_image_processor
            _lowerCAmelCase												=							prepare_img()
            _lowerCAmelCase												=							image_processor(images=a__      ,      return_tensors="""pt"""  ).to(a__  )

            # forward pass
            with torch.no_grad():
                  _lowerCAmelCase												=							model(**a__  )

            # verify the logits
            _lowerCAmelCase												=							torch.Size((1, 10_00)  )
            self.assertEqual(outputs.logits.shape      ,      a__  )

            _lowerCAmelCase												=							torch.tensor([-1.02_66, 0.19_12, -1.28_61]  ).to(a__  )

            self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,      a__      ,      atol=1E-4  )  )




      @slow
      @require_accelerate
      @require_torch_gpu
      def     lowercase__   (      self      :			List[str]  )    ->				Optional[Any]:
            _lowerCAmelCase												=							DeiTModel.from_pretrained(
                """facebook/deit-base-distilled-patch16-224"""      ,      torch_dtype=torch.floataa      ,      device_map="""auto"""  )
            _lowerCAmelCase												=							self.default_image_processor

            _lowerCAmelCase												=							prepare_img()
            _lowerCAmelCase												=							image_processor(images=a__      ,      return_tensors="""pt"""  )
            _lowerCAmelCase												=							inputs.pixel_values.to(a__  )

            # forward pass to make sure inference works in fp16
            with torch.no_grad():
                  _lowerCAmelCase												=							model(a__  )













'''simple docstring'''






import json
import logging
import os
import sys
from pathlib import Path

import finetune_rag

from transformers.file_utils import is_apex_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    require_ray,
    require_torch_gpu,
    require_torch_multi_gpu,
)


logging.basicConfig(level=logging.DEBUG)
A__    :       Any              =logging.getLogger()

A__    :       int              =logging.StreamHandler(sys.stdout)
logger.addHandler(stream_handler)




class 			UpperCAmelCase					(      snake_case_       ):




      def     lowercase__   (      self      :			Optional[Any]      ,      __snake_case      :			Any  )    ->				int:
            os.makedirs(__snake_case      ,      exist_ok=__snake_case  )
            _lowerCAmelCase												=							{"""source""": """What is love ?""", """target""": """life"""}
            _lowerCAmelCase												=							{"""train""": 12, """val""": 2, """test""": 2}
            for split in ["train", "test", "val"]:
                  for field in ["source", "target"]:
                        _lowerCAmelCase												=							"""\n""".join([contents[field]] * n_lines[split]  )
                        with open(os.path.join(__snake_case      ,      f"{split}.{field}"  )      ,      """w"""  ) as f:
                              f.write(__snake_case  )




      def     lowercase__   (      self      :			Tuple      ,      __snake_case      :			int      ,      __snake_case      :			str = "pytorch"  )    ->				int:
            _lowerCAmelCase												=							self.get_auto_remove_tmp_dir()
            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """output"""  )
            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """data"""  )
            self._create_dummy_data(data_dir=__snake_case  )

            _lowerCAmelCase												=							f"\n                --data_dir {data_dir} \\n                --output_dir {output_dir} \\n                --model_name_or_path facebook/rag-sequence-base \\n                --model_type rag_sequence \\n                --do_train \\n                --do_predict \\n                --n_val -1 \\n                --val_check_interval 1.0 \\n                --train_batch_size 2 \\n                --eval_batch_size 1 \\n                --max_source_length 25 \\n                --max_target_length 25 \\n                --val_max_target_length 25 \\n                --test_max_target_length 25 \\n                --label_smoothing 0.1 \\n                --dropout 0.1 \\n                --attention_dropout 0.1 \\n                --weight_decay 0.001 \\n                --adam_epsilon 1e-08 \\n                --max_grad_norm 0.1 \\n                --lr_scheduler polynomial \\n                --learning_rate 3e-04 \\n                --num_train_epochs 1 \\n                --warmup_steps 4 \\n                --gradient_accumulation_steps 1 \\n                --distributed-port 8787 \\n                --use_dummy_dataset 1 \\n                --distributed_retriever {distributed_retriever} \\n            ".split()

            if gpus > 0:
                  testargs.append(f"--gpus={gpus}"  )
                  if is_apex_available():
                        testargs.append("""--fp16"""  )
            else:
                  testargs.append("""--gpus=0"""  )
                  testargs.append("""--distributed_backend=ddp_cpu"""  )
                  testargs.append("""--num_processes=2"""  )

            _lowerCAmelCase												=							[sys.executable, str(Path(finetune_rag.__file__  ).resolve()  )] + testargs
            execute_subprocess_async(__snake_case      ,      env=self.get_env()  )

            _lowerCAmelCase												=							os.path.join(__snake_case      ,      """metrics.json"""  )
            with open(__snake_case  ) as f:
                  _lowerCAmelCase												=							json.load(__snake_case  )
            return result




      @require_torch_gpu
      def     lowercase__   (      self      :			Dict  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_multi_gpu
      def     lowercase__   (      self      :			int  )    ->				Dict:
            _lowerCAmelCase												=							self._run_finetune(gpus=2  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_gpu
      @require_ray
      def     lowercase__   (      self      :			Union[str, Any]  )    ->				Union[str, Any]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1      ,      distributed_retriever="""ray"""  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )




      @require_torch_multi_gpu
      @require_ray
      def     lowercase__   (      self      :			int  )    ->				List[str]:
            _lowerCAmelCase												=							self._run_finetune(gpus=1      ,      distributed_retriever="""ray"""  )
            self.assertGreaterEqual(result["""test"""][0]["""test_avg_em"""]      ,      0.2  )






