import inspect
import unittest

from transformers import DataaVecVisionConfig
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,
         DataaVecVisionForImageClassification,
         DataaVecVisionForSemanticSegmentation,
         DataaVecVisionModel,
     )
     from transformers.models.dataavec.modeling_dataavec_vision import DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
     from PIL import Image

     from transformers import BeitImageProcessor







class 			__snake_case   :
 def __init__(  self		,   snake_case__		,   snake_case__=100		,   snake_case__=13		,   snake_case__=30		,   snake_case__=2		,   snake_case__=3		,   snake_case__=True		,   snake_case__=True		,   snake_case__=32		,   snake_case__=4		,   snake_case__=4		,   snake_case__=37		,   snake_case__="gelu"		,   snake_case__=0.1		,   snake_case__=0.1		,   snake_case__=10		,   snake_case__=0.02		,   snake_case__=3		,   snake_case__=None		,   snake_case__=[0, 1, 2, 3]		,   )	->							int:
       '''simple docstring'''





       UpperCAmelCase    : Union[str, Any]      =parent
       UpperCAmelCase    : Dict      =100
       UpperCAmelCase    : Dict      =batch_size
       UpperCAmelCase    : Optional[int]      =image_size
       UpperCAmelCase    : Optional[int]      =patch_size
       UpperCAmelCase    : Dict      =num_channels
       UpperCAmelCase    : Dict      =is_training
       UpperCAmelCase    : Tuple      =use_labels
       UpperCAmelCase    : int      =hidden_size
       UpperCAmelCase    : Tuple      =num_hidden_layers
       UpperCAmelCase    : Dict      =num_attention_heads
       UpperCAmelCase    : int      =intermediate_size
       UpperCAmelCase    : str      =hidden_act
       UpperCAmelCase    : Union[str, Any]      =hidden_dropout_prob
       UpperCAmelCase    : List[Any]      =attention_probs_dropout_prob
       UpperCAmelCase    : Optional[Any]      =type_sequence_label_size
       UpperCAmelCase    : Optional[int]      =initializer_range
       UpperCAmelCase    : str      =scope
       UpperCAmelCase    : Optional[int]      =out_indices
       UpperCAmelCase    : Optional[int]      =num_labels

       # in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
       UpperCAmelCase    : Any      =(image_size // patch_size) ** 2
       UpperCAmelCase    : List[str]      =num_patches + 1
 def 	UpperCAmelCase__       (  self						)	->							str:
       '''simple docstring'''





       UpperCAmelCase    : List[Any]      =floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

       UpperCAmelCase    : Any      =self.get_config()

       return config, pixel_values, labels, pixel_labels
 def 	UpperCAmelCase__       (  self						)	->							str:
       '''simple docstring'''





       return DataaVecVisionConfig(
           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=snake_case__		,   initializer_range=self.initializer_range		,   out_indices=self.out_indices		,   )
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__						)	->							str:
       '''simple docstring'''





       UpperCAmelCase    : str      =DataaVecVisionModel(config=snake_case__						)
       model.to(snake_case__						)
       model.eval()
       UpperCAmelCase    : Optional[Any]      =model(snake_case__						)
       # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
       UpperCAmelCase    : Any      =(self.image_size // self.patch_size) ** 2
       self.parent.assertEqual(result.last_hidden_state.shape		,   (self.batch_size, num_patches + 1, self.hidden_size)						)
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__						)	->							Dict:
       '''simple docstring'''





       UpperCAmelCase    : List[Any]      =self.type_sequence_label_size
       UpperCAmelCase    : str      =DataaVecVisionForImageClassification(snake_case__						)
       model.to(snake_case__						)
       model.eval()
       UpperCAmelCase    : Optional[int]      =model(snake_case__		,   labels=snake_case__						)
       self.parent.assertEqual(result.logits.shape		,   (self.batch_size, self.type_sequence_label_size)						)
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__						)	->							str:
       '''simple docstring'''





       UpperCAmelCase    : Union[str, Any]      =self.num_labels
       UpperCAmelCase    : Optional[int]      =DataaVecVisionForSemanticSegmentation(snake_case__						)
       model.to(snake_case__						)
       model.eval()
       UpperCAmelCase    : Optional[int]      =model(snake_case__						)
       self.parent.assertEqual(
           result.logits.shape		,   (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)						)
       UpperCAmelCase    : Tuple      =model(snake_case__		,   labels=snake_case__						)
       self.parent.assertEqual(
           result.logits.shape		,   (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)						)






 def 	UpperCAmelCase__       (  self						)	->							Optional[int]:
       '''simple docstring'''





       UpperCAmelCase    : Dict      =self.prepare_config_and_inputs()
       UpperCAmelCase  ,						UpperCAmelCase  ,						UpperCAmelCase  ,						UpperCAmelCase    : Dict      =config_and_inputs
       UpperCAmelCase    : List[str]      ={'''pixel_values''': pixel_values}
       return config, inputs_dict







@require_torch
class 			__snake_case   (     lowerCamelCase__					,   lowerCamelCase__					,   unittest.TestCase   ):
 __lowerCamelCase    :   Any     =      (
     (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
     if is_torch_available()
     else ()
 )
 __lowerCamelCase    :   List[Any]     =      (
     {
         """feature-extraction""": DataaVecVisionModel,
         """image-classification""": DataaVecVisionForImageClassification,
         """image-segmentation""": DataaVecVisionForSemanticSegmentation,
     }
     if is_torch_available()
     else {}
 )

 __lowerCamelCase    :   str     =      False
 __lowerCamelCase    :   Optional[Any]     =      False
 __lowerCamelCase    :   List[str]     =      False
 def 	UpperCAmelCase__       (  self						)	->							Tuple:
       '''simple docstring'''





       UpperCAmelCase    : Tuple      =DataaVecVisionModelTester(self						)
       UpperCAmelCase    : List[str]      =ConfigTester(
           self		,   config_class=snake_case__		,   has_text_modality=snake_case__		,   hidden_size=37						)
 def 	UpperCAmelCase__       (  self						)	->							Dict:
       '''simple docstring'''





       self.config_tester.run_common_tests()
 def 	UpperCAmelCase__       (  self						)	->							List[str]:
       '''simple docstring'''





       pass
 @require_torch_multi_gpu
 @unittest.skip(
     reason='''Data2VecVision has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''						)
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       pass
 def 	UpperCAmelCase__       (  self						)	->							Optional[Any]:
       '''simple docstring'''





       UpperCAmelCase  ,						UpperCAmelCase    : Tuple      =self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             UpperCAmelCase    : Optional[int]      =model_class(snake_case__						)
             self.assertIsInstance(model.get_input_embeddings()		,   (nn.Module)						)
             UpperCAmelCase    : List[str]      =model.get_output_embeddings()
             self.assertTrue(x is None or isinstance(snake_case__		,   nn.Linear						)						)
 def 	UpperCAmelCase__       (  self						)	->							Tuple:
       '''simple docstring'''





       UpperCAmelCase  ,						UpperCAmelCase    : Optional[Any]      =self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             UpperCAmelCase    : Optional[Any]      =model_class(snake_case__						)
             UpperCAmelCase    : Dict      =inspect.signature(model.forward						)
             # signature.parameters is an OrderedDict => so arg_names order is deterministic
             UpperCAmelCase    : List[Any]      =[*signature.parameters.keys()]

             UpperCAmelCase    : str      =['''pixel_values''']
             self.assertListEqual(arg_names[:1]		,   snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       UpperCAmelCase    : Dict      =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							str:
       '''simple docstring'''





       UpperCAmelCase    : str      =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_segmentation(*snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       if not self.model_tester.is_training:
             return

       UpperCAmelCase  ,						UpperCAmelCase    : List[str]      =self.model_tester.prepare_config_and_inputs_for_common()
       UpperCAmelCase    : int      =True

       for model_class in self.all_model_classes:
             if model_class in [*get_values(snake_case__						)]:
                   continue

             UpperCAmelCase    : Dict      =model_class(snake_case__						)
             model.to(snake_case__						)
             model.train()
             UpperCAmelCase    : Optional[Any]      =self._prepare_for_class(snake_case__		,   snake_case__		,   return_labels=snake_case__						)
             UpperCAmelCase    : str      =model(**snake_case__						).loss
             loss.backward()
 def 	UpperCAmelCase__       (  self						)	->							str:
       '''simple docstring'''





       UpperCAmelCase  ,						UpperCAmelCase    : Union[str, Any]      =self.model_tester.prepare_config_and_inputs_for_common()
       if not self.model_tester.is_training:
             return

       UpperCAmelCase    : Dict      =False
       UpperCAmelCase    : Optional[Any]      =True

       for model_class in self.all_model_classes:
             if model_class in [*get_values(snake_case__						)] or not model_class.supports_gradient_checkpointing:
                   continue
             # TODO: remove the following 3 lines once we have a MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING
             # this can then be incorporated into _prepare_for_class in test_modeling_common.py
             elif model_class.__name__ == "Data2VecVisionForSemanticSegmentation":
                   UpperCAmelCase  ,						UpperCAmelCase  ,						UpperCAmelCase  ,						UpperCAmelCase    : Any      =inputs_dict['''pixel_values'''].shape
                   UpperCAmelCase    : Dict      =torch.zeros(
                       [self.model_tester.batch_size, height, width]		,   device=snake_case__						).long()
             UpperCAmelCase    : Dict      =model_class(snake_case__						)
             model.gradient_checkpointing_enable()
             model.to(snake_case__						)
             model.train()
             UpperCAmelCase    : Tuple      =self._prepare_for_class(snake_case__		,   snake_case__		,   return_labels=snake_case__						)
             UpperCAmelCase    : List[str]      =model(**snake_case__						).loss
             loss.backward()
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       UpperCAmelCase  ,						UpperCAmelCase    : Tuple      =self.model_tester.prepare_config_and_inputs_for_common()

       UpperCAmelCase    : Union[str, Any]      =_config_zero_init(snake_case__						)
       for model_class in self.all_model_classes:
             UpperCAmelCase    : Dict      =model_class(config=snake_case__						)
             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'''		,   )
 def 	UpperCAmelCase__       (  self		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__=2e-4		,   snake_case__="outputs"		,   snake_case__=None						)	->							Optional[Any]:
       '''simple docstring'''





       super().check_pt_tf_outputs(snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__						)
 def 	UpperCAmelCase__       (  self						)	->							int:
       '''simple docstring'''





       UpperCAmelCase    : Optional[Any]      =self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*snake_case__						)






 @slow
 def 	UpperCAmelCase__       (  self						)	->							List[str]:
       '''simple docstring'''





       for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             UpperCAmelCase    : Tuple      =DataaVecVisionModel.from_pretrained(snake_case__						)
             self.assertIsNotNone(snake_case__						)





def   lowerCAmelCase_ (				)->  Dict:






      '''simple docstring'''





      UpperCAmelCase    : Any      =Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
      return image







@require_torch
@require_vision
class 			__snake_case   (     unittest.TestCase   ):
 @cached_property
 def 	UpperCAmelCase__       (  self						)	->							List[Any]:
       '''simple docstring'''





       return (
           BeitImageProcessor.from_pretrained('''facebook/data2vec-vision-base-ft1k'''						) if is_vision_available() else None
       )






 @slow
 def 	UpperCAmelCase__       (  self						)	->							Union[str, Any]:
       '''simple docstring'''





       UpperCAmelCase    : List[Any]      =DataaVecVisionForImageClassification.from_pretrained('''facebook/data2vec-vision-base-ft1k'''						).to(
           snake_case__						)

       UpperCAmelCase    : Optional[int]      =self.default_image_processor
       UpperCAmelCase    : Tuple      =prepare_img()
       UpperCAmelCase    : List[Any]      =image_processor(images=snake_case__		,   return_tensors='''pt'''						).to(snake_case__						)

       # forward pass
       with torch.no_grad():
             UpperCAmelCase    : str      =model(**snake_case__						)
       UpperCAmelCase    : int      =outputs.logits

       # verify the logits
       UpperCAmelCase    : Tuple      =torch.Size((1, 1000)						)
       self.assertEqual(logits.shape		,   snake_case__						)

       UpperCAmelCase    : List[str]      =torch.tensor([0.3277, -0.1395, 0.0911]						).to(snake_case__						)

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

       UpperCAmelCase    : Dict      =[model.config.labelaid[i] for i in ['''remote control, remote''', '''tabby, tabby cat''']]
       self.assertEqual(logits[0].topk(2						).indices.cpu().tolist()		,   snake_case__						)

import os
from argparse import ArgumentParser
from pathlib import Path
from shutil import copyfile
from typing import List

from datasets import config
from datasets.builder import DatasetBuilder
from datasets.commands import BaseDatasetsCLICommand
from datasets.download.download_config import DownloadConfig
from datasets.download.download_manager import DownloadMode
from datasets.load import dataset_module_factory, import_main_class
from datasets.utils.info_utils import VerificationMode





def   lowerCAmelCase_ (				__lowerCAmelCase		,					**__lowerCAmelCase						)->  Union[str, Any]:






      '''simple docstring'''





      return RunBeamCommand(
          args.dataset		,					args.name		,					args.cache_dir		,					args.beam_pipeline_options		,					args.data_dir		,					args.all_configs		,					args.save_info or args.save_infos		,					args.ignore_verifications		,					args.force_redownload		,					**__lowerCAmelCase		,					)







class 			__snake_case   (     lowerCamelCase__   ):
 @staticmethod
 def 	UpperCAmelCase__       (  snake_case__						)	->							List[str]:
       '''simple docstring'''





       UpperCAmelCase    : List[str]      =parser.add_parser('''run_beam'''		,   help='''Run a Beam dataset processing pipeline'''						)
       run_beam_parser.add_argument('''dataset'''		,   type=snake_case__		,   help='''Name of the dataset to download'''						)
       run_beam_parser.add_argument('''--name'''		,   type=snake_case__		,   default=snake_case__		,   help='''Dataset config name'''						)
       run_beam_parser.add_argument(
           '''--cache_dir'''		,   type=snake_case__		,   default=snake_case__		,   help='''Cache directory where the datasets are stored'''		,   )
       run_beam_parser.add_argument(
           '''--beam_pipeline_options'''		,   type=snake_case__		,   default=''''''		,   help='''Beam pipeline options, separated by commas. Example:: `--beam_pipeline_options=job_name=my-job,project=my-project`'''		,   )
       run_beam_parser.add_argument(
           '''--data_dir'''		,   type=snake_case__		,   default=snake_case__		,   help='''Can be used to specify a manual directory to get the files from'''		,   )
       run_beam_parser.add_argument('''--all_configs'''		,   action='''store_true'''		,   help='''Test all dataset configurations'''						)
       run_beam_parser.add_argument('''--save_info'''		,   action='''store_true'''		,   help='''Save the dataset infos file'''						)
       run_beam_parser.add_argument(
           '''--ignore_verifications'''		,   action='''store_true'''		,   help='''Run the test without checksums and splits checks'''						)
       run_beam_parser.add_argument('''--force_redownload'''		,   action='''store_true'''		,   help='''Force dataset redownload'''						)
       # aliases
       run_beam_parser.add_argument('''--save_infos'''		,   action='''store_true'''		,   help='''alias for save_info'''						)
       run_beam_parser.set_defaults(func=snake_case__						)
 def __init__(  self		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   snake_case__		,   **snake_case__		,   )	->							Any:
       '''simple docstring'''





       UpperCAmelCase    : str      =dataset
       UpperCAmelCase    : Optional[Any]      =name
       UpperCAmelCase    : str      =cache_dir
       UpperCAmelCase    : Union[str, Any]      =beam_pipeline_options
       UpperCAmelCase    : List[str]      =data_dir
       UpperCAmelCase    : int      =all_configs
       UpperCAmelCase    : List[Any]      =save_infos
       UpperCAmelCase    : Dict      =ignore_verifications
       UpperCAmelCase    : Optional[int]      =force_redownload
       UpperCAmelCase    : Optional[int]      =config_kwargs






 def 	UpperCAmelCase__       (  self						)	->							str:
       '''simple docstring'''





       import apache_beam as beam

       if self._name is not None and self._all_configs:
             print('''Both parameters `name` and `all_configs` can\'t be used at once.'''						)
             exit(1						)
       UpperCAmelCase  ,						UpperCAmelCase    : str      =self._dataset, self._name
       UpperCAmelCase    : Any      =dataset_module_factory(snake_case__						)
       UpperCAmelCase    : str      =import_main_class(dataset_module.module_path						)
       UpperCAmelCase    : List[DatasetBuilder]      =[]
       if self._beam_pipeline_options:
             UpperCAmelCase    : Optional[Any]      =beam.options.pipeline_options.PipelineOptions(
                 flags=[f'''--{opt.strip()}''' for opt in self._beam_pipeline_options.split(''','''						) if opt]						)
       else:
             UpperCAmelCase    : int      =None
       if self._all_configs and len(builder_cls.BUILDER_CONFIGS						) > 0:
             for builder_config in builder_cls.BUILDER_CONFIGS:
                   builders.append(
                       builder_cls(
                           config_name=builder_config.name		,   data_dir=self._data_dir		,   hash=dataset_module.hash		,   beam_options=snake_case__		,   cache_dir=self._cache_dir		,   base_path=dataset_module.builder_kwargs.get('''base_path'''						)		,   )						)
       else:
             builders.append(
                 builder_cls(
                     config_name=snake_case__		,   data_dir=self._data_dir		,   beam_options=snake_case__		,   cache_dir=self._cache_dir		,   base_path=dataset_module.builder_kwargs.get('''base_path'''						)		,   **self._config_kwargs		,   )						)

       for builder in builders:
             builder.download_and_prepare(
                 download_mode=DownloadMode.REUSE_CACHE_IF_EXISTS
                 if not self._force_redownload
                 else DownloadMode.FORCE_REDOWNLOAD		,   download_config=DownloadConfig(cache_dir=config.DOWNLOADED_DATASETS_PATH						)		,   verification_mode=VerificationMode.NO_CHECKS
                 if self._ignore_verifications
                 else VerificationMode.ALL_CHECKS		,   try_from_hf_gcs=snake_case__		,   )
             if self._save_infos:
                   builder._save_infos()

       print('''Apache beam run successful.'''						)

       # If save_infos=True, the dataset infos file is created next to the loaded module file.
       # Let's move it to the original directory of the dataset script, to allow the user to
       # upload them on S3 at the same time afterwards.
       if self._save_infos:
             UpperCAmelCase    : Dict      =os.path.join(builder_cls.get_imported_module_dir()		,   config.DATASETDICT_INFOS_FILENAME						)

             UpperCAmelCase    : List[str]      =Path(snake_case__						).name + '''.py'''

             UpperCAmelCase    : Tuple      =os.path.join(snake_case__		,   snake_case__						)
             if os.path.isfile(snake_case__						):
                   UpperCAmelCase    : Dict      =os.path.dirname(snake_case__						)
             elif os.path.isfile(snake_case__						):
                   UpperCAmelCase    : Union[str, Any]      =path
             else:  # in case of a remote dataset
                   print(f'''Dataset Infos file saved at {dataset_infos_path}'''						)
                   exit(1						)

             # Move datasetinfo back to the user
             UpperCAmelCase    : Tuple      =os.path.join(snake_case__		,   config.DATASETDICT_INFOS_FILENAME						)
             copyfile(snake_case__		,   snake_case__						)
             print(f'''Dataset Infos file saved at {user_dataset_infos_path}'''						)
