


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       :						Any	,      _lowerCAmelCase       :						List[str]	,      _lowerCAmelCase       :						str=1_0_0	,      _lowerCAmelCase       :						Optional[int]=1_3	,      _lowerCAmelCase       :						Union[str, Any]=3_0	,      _lowerCAmelCase       :						Tuple=2	,      _lowerCAmelCase       :						Dict=3	,      _lowerCAmelCase       :						Any=True	,      _lowerCAmelCase       :						Any=True	,      _lowerCAmelCase       :						int=3_2	,      _lowerCAmelCase       :						Dict=4	,      _lowerCAmelCase       :						Tuple=4	,      _lowerCAmelCase       :						Dict=3_7	,      _lowerCAmelCase       :						Any="gelu"	,      _lowerCAmelCase       :						Dict=0.1	,      _lowerCAmelCase       :						List[Any]=0.1	,      _lowerCAmelCase       :						Optional[Any]=1_0	,      _lowerCAmelCase       :						Optional[Any]=0.02	,      _lowerCAmelCase       :						Tuple=3	,      _lowerCAmelCase       :						Dict=None	,      _lowerCAmelCase       :						Optional[int]=[0, 1, 2, 3]	,      )							->					Union[str, Any]:



											'''simple docstring'''





											__magic_name__					=      parent
											__magic_name__					=      1_0_0
											__magic_name__					=      batch_size
											__magic_name__					=      image_size
											__magic_name__					=      patch_size
											__magic_name__					=      num_channels
											__magic_name__					=      is_training
											__magic_name__					=      use_labels
											__magic_name__					=      hidden_size
											__magic_name__					=      num_hidden_layers
											__magic_name__					=      num_attention_heads
											__magic_name__					=      intermediate_size
											__magic_name__					=      hidden_act
											__magic_name__					=      hidden_dropout_prob
											__magic_name__					=      attention_probs_dropout_prob
											__magic_name__					=      type_sequence_label_size
											__magic_name__					=      initializer_range
											__magic_name__					=      scope
											__magic_name__					=      out_indices
											__magic_name__					=      num_labels

											# in BeiT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
											__magic_name__					=      (image_size // patch_size) ** 2
											__magic_name__					=      num_patches + 1
					def 	a			(						self       :						str  )							->					Dict:



											'''simple docstring'''





											__magic_name__					=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

											__magic_name__					=      None
											__magic_name__					=      None
											if self.use_labels:
																	__magic_name__					=      ids_tensor([self.batch_size]	,      self.type_sequence_label_size  )
																	__magic_name__					=      ids_tensor([self.batch_size, self.image_size, self.image_size]	,      self.num_labels  )

											__magic_name__					=      self.get_config()

											return config, pixel_values, labels, pixel_labels
					def 	a			(						self       :						Tuple  )							->					Union[str, Any]:



											'''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=_lowerCAmelCase	,      initializer_range=self.initializer_range	,      out_indices=self.out_indices	,      )
					def 	a			(						self       :						Optional[Any]	,      _lowerCAmelCase       :						Optional[int]	,      _lowerCAmelCase       :						Dict	,      _lowerCAmelCase       :						Any	,      _lowerCAmelCase       :						Tuple  )							->					Any:



											'''simple docstring'''





											__magic_name__					=      DataaVecVisionModel(config=_lowerCAmelCase  )
											model.to(_lowerCAmelCase  )
											model.eval()
											__magic_name__					=      model(_lowerCAmelCase  )
											# expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
											__magic_name__					=      (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 	a			(						self       :						int	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						List[str]	,      _lowerCAmelCase       :						Union[str, Any]	,      _lowerCAmelCase       :						int  )							->					Dict:



											'''simple docstring'''





											__magic_name__					=      self.type_sequence_label_size
											__magic_name__					=      DataaVecVisionForImageClassification(_lowerCAmelCase  )
											model.to(_lowerCAmelCase  )
											model.eval()
											__magic_name__					=      model(_lowerCAmelCase	,      labels=_lowerCAmelCase  )
											self.parent.assertEqual(result.logits.shape	,      (self.batch_size, self.type_sequence_label_size)  )
					def 	a			(						self       :						Tuple	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						List[Any]	,      _lowerCAmelCase       :						List[Any]  )							->					Tuple:



											'''simple docstring'''





											__magic_name__					=      self.num_labels
											__magic_name__					=      DataaVecVisionForSemanticSegmentation(_lowerCAmelCase  )
											model.to(_lowerCAmelCase  )
											model.eval()
											__magic_name__					=      model(_lowerCAmelCase  )
											self.parent.assertEqual(
											    result.logits.shape	,      (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )
											__magic_name__					=      model(_lowerCAmelCase	,      labels=_lowerCAmelCase  )
											self.parent.assertEqual(
											    result.logits.shape	,      (self.batch_size, self.num_labels, self.image_size * 2, self.image_size * 2)  )







					def 	a			(						self       :						List[Any]  )							->					Tuple:



											'''simple docstring'''





											__magic_name__					=      self.prepare_config_and_inputs()
											__magic_name__		,      __magic_name__		,      __magic_name__		,      __magic_name__					=      config_and_inputs
											__magic_name__					=      {'pixel_values': pixel_values}
											return config, inputs_dict







@require_torch
class      __snake_case				(   _UpperCAmelCase	,      _UpperCAmelCase	,      unittest.TestCase):
					SCREAMING_SNAKE_CASE_		:       Union[str, Any]								=			(
					    (DataaVecVisionModel, DataaVecVisionForImageClassification, DataaVecVisionForSemanticSegmentation)
					    if is_torch_available()
					    else ()
					)
					SCREAMING_SNAKE_CASE_		:       Optional[int]								=			(
					    {
					        "feature-extraction": DataaVecVisionModel,
					        "image-classification": DataaVecVisionForImageClassification,
					        "image-segmentation": DataaVecVisionForSemanticSegmentation,
					    }
					    if is_torch_available()
					    else {}
					)

					SCREAMING_SNAKE_CASE_		:       Optional[int]								=			False
					SCREAMING_SNAKE_CASE_		:       Optional[int]								=			False
					SCREAMING_SNAKE_CASE_		:       List[str]								=			False
					def 	a			(						self       :						int  )							->					Union[str, Any]:



											'''simple docstring'''





											__magic_name__					=      DataaVecVisionModelTester(self  )
											__magic_name__					=      ConfigTester(
											    self	,      config_class=_lowerCAmelCase	,      has_text_modality=_lowerCAmelCase	,      hidden_size=3_7  )
					def 	a			(						self       :						str  )							->					Any:



											'''simple docstring'''





											self.config_tester.run_common_tests()
					def 	a			(						self       :						Optional[int]  )							->					int:



											'''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 	a			(						self       :						Optional[Any]  )							->					Union[str, Any]:



											'''simple docstring'''





											pass
					def 	a			(						self       :						int  )							->					str:



											'''simple docstring'''





											__magic_name__		,      __magic_name__					=      self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
																	__magic_name__					=      model_class(_lowerCAmelCase  )
																	self.assertIsInstance(model.get_input_embeddings()	,      (nn.Module)  )
																	__magic_name__					=      model.get_output_embeddings()
																	self.assertTrue(x is None or isinstance(_lowerCAmelCase	,      nn.Linear  )  )
					def 	a			(						self       :						int  )							->					Any:



											'''simple docstring'''





											__magic_name__		,      __magic_name__					=      self.model_tester.prepare_config_and_inputs_for_common()

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

																	__magic_name__					=      ['pixel_values']
																	self.assertListEqual(arg_names[:1]	,      _lowerCAmelCase  )
					def 	a			(						self       :						str  )							->					Dict:



											'''simple docstring'''





											__magic_name__					=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*_lowerCAmelCase  )
					def 	a			(						self       :						List[Any]  )							->					Optional[Any]:



											'''simple docstring'''





											__magic_name__					=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_segmentation(*_lowerCAmelCase  )
					def 	a			(						self       :						List[str]  )							->					Union[str, Any]:



											'''simple docstring'''





											if not self.model_tester.is_training:
																	return

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

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

																	__magic_name__					=      model_class(_lowerCAmelCase  )
																	model.to(_lowerCAmelCase  )
																	model.train()
																	__magic_name__					=      self._prepare_for_class(_lowerCAmelCase	,      _lowerCAmelCase	,      return_labels=_lowerCAmelCase  )
																	__magic_name__					=      model(**_lowerCAmelCase  ).loss
																	loss.backward()
					def 	a			(						self       :						Union[str, Any]  )							->					Any:



											'''simple docstring'''





											__magic_name__		,      __magic_name__					=      self.model_tester.prepare_config_and_inputs_for_common()
											if not self.model_tester.is_training:
																	return

											__magic_name__					=      False
											__magic_name__					=      True

											for model_class in self.all_model_classes:
																	if model_class in [*get_values(_lowerCAmelCase  )] 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":
																							__magic_name__		,      __magic_name__		,      __magic_name__		,      __magic_name__					=      inputs_dict['pixel_values'].shape
																							__magic_name__					=      torch.zeros(
																							    [self.model_tester.batch_size, height, width]	,      device=_lowerCAmelCase  ).long()
																	__magic_name__					=      model_class(_lowerCAmelCase  )
																	model.gradient_checkpointing_enable()
																	model.to(_lowerCAmelCase  )
																	model.train()
																	__magic_name__					=      self._prepare_for_class(_lowerCAmelCase	,      _lowerCAmelCase	,      return_labels=_lowerCAmelCase  )
																	__magic_name__					=      model(**_lowerCAmelCase  ).loss
																	loss.backward()
					def 	a			(						self       :						Tuple  )							->					int:



											'''simple docstring'''





											__magic_name__		,      __magic_name__					=      self.model_tester.prepare_config_and_inputs_for_common()

											__magic_name__					=      _config_zero_init(_lowerCAmelCase  )
											for model_class in self.all_model_classes:
																	__magic_name__					=      model_class(config=_lowerCAmelCase  )
																	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 	a			(						self       :						Optional[Any]	,      _lowerCAmelCase       :						List[Any]	,      _lowerCAmelCase       :						List[Any]	,      _lowerCAmelCase       :						Optional[Any]	,      _lowerCAmelCase       :						str=2E-4	,      _lowerCAmelCase       :						Union[str, Any]="outputs"	,      _lowerCAmelCase       :						Union[str, Any]=None  )							->					Union[str, Any]:



											'''simple docstring'''





											super().check_pt_tf_outputs(_lowerCAmelCase	,      _lowerCAmelCase	,      _lowerCAmelCase	,      _lowerCAmelCase	,      _lowerCAmelCase	,      _lowerCAmelCase  )
					def 	a			(						self       :						Dict  )							->					Tuple:



											'''simple docstring'''





											__magic_name__					=      self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*_lowerCAmelCase  )







					@slow
					def 	a			(						self       :						List[str]  )							->					List[Any]:



											'''simple docstring'''





											for model_name in DATA2VEC_VISION_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	__magic_name__					=      DataaVecVisionModel.from_pretrained(_lowerCAmelCase  )
																	self.assertIsNotNone(_lowerCAmelCase  )







def 				UpperCamelCase_	(					)->   Any:




						'''simple docstring'''
						__magic_name__					=      Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'  )
						return image




@require_torch
@require_vision
class      __snake_case				(   unittest.TestCase):
					@cached_property
					def 	a			(						self       :						Optional[int]  )							->					Tuple:



											'''simple docstring'''





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







					@slow
					def 	a			(						self       :						int  )							->					str:



											'''simple docstring'''





											__magic_name__					=      DataaVecVisionForImageClassification.from_pretrained('facebook/data2vec-vision-base-ft1k'  ).to(
											    _lowerCAmelCase  )

											__magic_name__					=      self.default_image_processor
											__magic_name__					=      prepare_img()
											__magic_name__					=      image_processor(images=_lowerCAmelCase	,      return_tensors='pt'  ).to(_lowerCAmelCase  )

											# forward pass
											with torch.no_grad():
																	__magic_name__					=      model(**_lowerCAmelCase  )
											__magic_name__					=      outputs.logits

											# verify the logits
											__magic_name__					=      torch.Size((1, 1_0_0_0)  )
											self.assertEqual(logits.shape	,      _lowerCAmelCase  )

											__magic_name__					=      torch.tensor([0.3277, -0.1395, 0.0911]  ).to(_lowerCAmelCase  )

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

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





import argparse
import logging

import datasets
from tokenizers import Tokenizer, decoders, normalizers, pre_tokenizers, processors
from tokenizers.models import Unigram
from tokenizers.trainers import UnigramTrainer

from transformers import AlbertTokenizerFast


UpperCamelCase			:					List[str]         =							logging.getLogger(__name__)
def 				UpperCamelCase_	(					)->   Dict:




						'''simple docstring'''
						__magic_name__					=      argparse.ArgumentParser(description='Train a unigram tokenizer on the wikitext dataset.'  )
						parser.add_argument(
						    '--dataset_name'						,      type=__snake_case						,      default='wikitext'						,      help='Name of the training. Explore datasets at: hf.co/datasets.'						,      )
						parser.add_argument(
						    '--dataset_config'						,      type=__snake_case						,      default='wikitext-103-raw-v1'						,      help='Configuration name of the dataset.'  )
						parser.add_argument(
						    '--batch_size'						,      type=__snake_case						,      default=10_00						,      help='Batch size during training.'						,      )
						parser.add_argument(
						    '--vocab_size'						,      type=__snake_case						,      default=1_00_48						,      help='Size of the desired vocabulary.'						,      )
						parser.add_argument(
						    '--limit'						,      default=__snake_case						,      type=__snake_case						,      help='Limit the number of shards (used for debugging).'						,      )
						parser.add_argument(
						    '--export_to_hub'						,      action='store_true'						,      )

						__magic_name__					=      parser.parse_args()
						return args
def 				UpperCamelCase_	(					__snake_case      :    Any  )->   Optional[int]:




						'''simple docstring'''
						__magic_name__					=      datasets.load_dataset(args.dataset_name						,      args.dataset_config						,      split='train'  )

						if args.limit is not None:
												__magic_name__					=      min(len(__snake_case  )						,      args.limit  )
												__magic_name__					=      dataset.select(range(__snake_case  )  )
												logger.info(F'''Limiting the dataset to {args.limit} entries.'''  )

						def batch_iterator():
												for i in range(0						,      len(__snake_case  )						,      args.batch_size  ):
																		yield dataset[i : i + args.batch_size]["text"]

    # Prepare the tokenizer.
						__magic_name__					=      Tokenizer(Unigram()  )
						__magic_name__					=      normalizers.Sequence([normalizers.Replace('``'						,      '"'  ), normalizers.Replace('\'\''						,      '"'  )]  )
						__magic_name__					=      pre_tokenizers.Metaspace()

						# Prepare the trainer.
						__magic_name__					=      UnigramTrainer(
						    unk_token='<unk>'						,      special_tokens=['[CLS]', '[SEP]', '<unk>', '<pad>', '[MASK]']						,      vocab_size=args.vocab_size						,      )

						logger.info('Training the tokenizer.'  )
						tokenizer.train_from_iterator(batch_iterator()						,      trainer=__snake_case  )
						logger.info('Tokenizer training complete!'  )

						__magic_name__					=      tokenizer.token_to_id('[CLS]'  )
						__magic_name__					=      tokenizer.token_to_id('[SEP]'  )
						__magic_name__					=      processors.TemplateProcessing(
						    single='[CLS]:0 $A:0 [SEP]:0'						,      pair='[CLS]:0 $A:0 [SEP]:0 $B:1 [SEP]:1'						,      special_tokens=[
						        ('[CLS]', cls_token_id),
						        ('[SEP]', sep_token_id),
						    ]						,      )
						__magic_name__					=      decoders.Metaspace()

						if args.export_to_hub:
												logger.info('Exporting the trained tokenzier to Hub.'  )
												__magic_name__					=      AlbertTokenizerFast(tokenizer_object=__snake_case  )
												new_tokenizer.push_to_hub('unigram-tokenizer-dataset'  )


if __name__ == "__main__":
					UpperCamelCase			:					Dict         =							parse_args()
					main(args)

