





import inspect
import unittest

from transformers import MobileNetVaConfig
from transformers.testing_utils import require_torch, 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 transformers import MobileNetVaForImageClassification, MobileNetVaModel
	from transformers.models.mobilenet_va.modeling_mobilenet_va import MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import MobileNetVaImageProcessor




class   UpperCamelCase_       (     UpperCamelCase__    ):


							def      _snake_case  (							self    :Union[str, Any]		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.config_class(**self.inputs_dict		)
									self.parent.assertTrue(hasattr(__A							,				"""tf_padding"""		)		)
									self.parent.assertTrue(hasattr(__A							,				"""depth_multiplier"""		)		)




class   UpperCamelCase_       :


							def __init__(							self    :List[str]							,				__A    :Optional[int]							,				__A    :int=13							,				__A    :Union[str, Any]=3							,				__A    :Tuple=32							,				__A    :Union[str, Any]=0.2_5							,				__A    :List[Any]=8							,				__A    :Dict=True							,				__A    :List[Any]=1024							,				__A    :Dict=32							,				__A    :Optional[int]="relu6"							,				__A    :Optional[Any]=0.1							,				__A    :Any=0.0_2							,				__A    :Tuple=True							,				__A    :int=True							,				__A    :Optional[int]=10							,				__A    :List[Any]=None							,				)   ->   Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      parent
									SCREAMING_SNAKE_CASE__			=      batch_size
									SCREAMING_SNAKE_CASE__			=      num_channels
									SCREAMING_SNAKE_CASE__			=      image_size
									SCREAMING_SNAKE_CASE__			=      depth_multiplier
									SCREAMING_SNAKE_CASE__			=      min_depth
									SCREAMING_SNAKE_CASE__			=      tf_padding
									SCREAMING_SNAKE_CASE__			=      int(last_hidden_size * depth_multiplier		)
									SCREAMING_SNAKE_CASE__			=      output_stride
									SCREAMING_SNAKE_CASE__			=      hidden_act
									SCREAMING_SNAKE_CASE__			=      classifier_dropout_prob
									SCREAMING_SNAKE_CASE__			=      use_labels
									SCREAMING_SNAKE_CASE__			=      is_training
									SCREAMING_SNAKE_CASE__			=      num_labels
									SCREAMING_SNAKE_CASE__			=      initializer_range
									SCREAMING_SNAKE_CASE__			=      scope


							def      _snake_case  (							self    :Union[str, Any]		)   ->   Union[str, Any]:



									"""simple docstring"""



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

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

									SCREAMING_SNAKE_CASE__			=      self.get_config()

									return config, pixel_values, labels, pixel_labels


							def      _snake_case  (							self    :Dict		)   ->   int:



									"""simple docstring"""



									return MobileNetVaConfig(
									    num_channels=self.num_channels							,				image_size=self.image_size							,				depth_multiplier=self.depth_multiplier							,				min_depth=self.min_depth							,				tf_padding=self.tf_padding							,				hidden_act=self.hidden_act							,				classifier_dropout_prob=self.classifier_dropout_prob							,				initializer_range=self.initializer_range							,				)


							def      _snake_case  (							self    :Optional[int]							,				__A    :Any							,				__A    :Optional[int]							,				__A    :int							,				__A    :Dict		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileNetVaModel(config=__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__			=      model(__A		)
									self.parent.assertEqual(
									    result.last_hidden_state.shape							,				(
									        self.batch_size,
									        self.last_hidden_size,
									        self.image_size // self.output_stride,
									        self.image_size // self.output_stride,
									    )							,				)


							def      _snake_case  (							self    :List[str]							,				__A    :Optional[int]							,				__A    :int							,				__A    :Optional[int]							,				__A    :str		)   ->   str:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.num_labels
									SCREAMING_SNAKE_CASE__			=      MobileNetVaForImageClassification(__A		)
									model.to(__A		)
									model.eval()
									SCREAMING_SNAKE_CASE__			=      model(__A							,				labels=__A		)
									self.parent.assertEqual(result.logits.shape							,				(self.batch_size, self.num_labels)		)


							def      _snake_case  (							self    :int		)   ->   Tuple:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      config_and_inputs
									SCREAMING_SNAKE_CASE__			=      {"""pixel_values""": pixel_values}
									return config, inputs_dict




@require_torch
class   UpperCamelCase_       (     UpperCamelCase__				,  UpperCamelCase__				,  unittest.TestCase    ):
							lowerCamelCase_           =					(MobileNetVaModel, MobileNetVaForImageClassification) if is_torch_available() else ()
							lowerCamelCase_           =					(
							    {"feature-extraction": MobileNetVaModel, "image-classification": MobileNetVaForImageClassification}
							    if is_torch_available()
							    else {}
							)

							lowerCamelCase_           =					False
							lowerCamelCase_           =					False
							lowerCamelCase_           =					False
							lowerCamelCase_           =					False


							def      _snake_case  (							self    :Dict		)   ->   List[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileNetVaModelTester(self		)
									SCREAMING_SNAKE_CASE__			=      MobileNetVaConfigTester(self							,				config_class=__A							,				has_text_modality=__A		)


							def      _snake_case  (							self    :List[str]		)   ->   List[Any]:



									"""simple docstring"""



									self.config_tester.run_common_tests()


							@unittest.skip(reason="""MobileNetV1 does not use inputs_embeds"""		)
							def      _snake_case  (							self    :Dict		)   ->   Optional[int]:



									"""simple docstring"""



									pass


							@unittest.skip(reason="""MobileNetV1 does not support input and output embeddings"""		)
							def      _snake_case  (							self    :Optional[int]		)   ->   Optional[Any]:



									"""simple docstring"""



									pass


							@unittest.skip(reason="""MobileNetV1 does not output attentions"""		)
							def      _snake_case  (							self    :Any		)   ->   int:



									"""simple docstring"""



									pass


							def      _snake_case  (							self    :List[Any]		)   ->   Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs_for_common()

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

											SCREAMING_SNAKE_CASE__			=      ["""pixel_values"""]
											self.assertListEqual(arg_names[:1]							,				__A		)


							def      _snake_case  (							self    :int		)   ->   Union[str, Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*__A		)


							def      _snake_case  (							self    :Optional[Any]		)   ->   Tuple:



									"""simple docstring"""



									def check_hidden_states_output(__A    :Optional[Any]							,				__A    :Optional[int]							,				__A    :List[str]		):
											SCREAMING_SNAKE_CASE__			=      model_class(__A		)
											model.to(__A		)
											model.eval()

											with torch.no_grad():
													SCREAMING_SNAKE_CASE__			=      model(**self._prepare_for_class(__A							,				__A		)		)

											SCREAMING_SNAKE_CASE__			=      outputs.hidden_states

											SCREAMING_SNAKE_CASE__			=      26
											self.assertEqual(len(__A		)							,				__A		)

									SCREAMING_SNAKE_CASE__	,						SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs_for_common()

									for model_class in self.all_model_classes:
											SCREAMING_SNAKE_CASE__			=      True
											check_hidden_states_output(__A							,				__A							,				__A		)

											# check that output_hidden_states also work using config
											del inputs_dict["output_hidden_states"]
											SCREAMING_SNAKE_CASE__			=      True

											check_hidden_states_output(__A							,				__A							,				__A		)


							def      _snake_case  (							self    :Any		)   ->   int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_image_classification(*__A		)


							@slow
							def      _snake_case  (							self    :Any		)   ->   Union[str, Any]:



									"""simple docstring"""



									for model_name in MOBILENET_V1_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											SCREAMING_SNAKE_CASE__			=      MobileNetVaModel.from_pretrained(__A		)
											self.assertIsNotNone(__A		)






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




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


							@cached_property
							def      _snake_case  (							self    :Any		)   ->   Union[str, Any]:



									"""simple docstring"""



									return (
									    MobileNetVaImageProcessor.from_pretrained("""google/mobilenet_v1_1.0_224"""		) if is_vision_available() else None
									)


							@slow
							def      _snake_case  (							self    :Dict		)   ->   Optional[int]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__			=      MobileNetVaForImageClassification.from_pretrained("""google/mobilenet_v1_1.0_224"""		).to(__A		)

									SCREAMING_SNAKE_CASE__			=      self.default_image_processor
									SCREAMING_SNAKE_CASE__			=      prepare_img()
									SCREAMING_SNAKE_CASE__			=      image_processor(images=__A							,				return_tensors="""pt"""		).to(__A		)

									# forward pass
									with torch.no_grad():
											SCREAMING_SNAKE_CASE__			=      model(**__A		)

									# verify the logits
									SCREAMING_SNAKE_CASE__			=      torch.Size((1, 1001)		)
									self.assertEqual(outputs.logits.shape							,				__A		)

									SCREAMING_SNAKE_CASE__			=      torch.tensor([-4.1_7_3_9, -1.1_2_3_3, 3.1_2_0_5]		).to(__A		)

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






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


_lowerCamelCase   				=			logging.getLogger(__name__)






def 				SCREAMING_SNAKE_CASE__							(       ):
		SCREAMING_SNAKE_CASE__			=      argparse.ArgumentParser(
		    description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""    )
		parser.add_argument(
		    """--dataset_name"""					,   type=UpperCamelCase__					,   default="""wikitext"""					,   help="""Name of the training. Explore datasets at: hf.co/datasets."""					,   )
		parser.add_argument(
		    """--dataset_config"""					,   type=UpperCamelCase__					,   default="""wikitext-103-raw-v1"""					,   help="""Configuration name of the dataset."""    )
		parser.add_argument(
		    """--tokenizer_name_or_path"""					,   type=UpperCamelCase__					,   default="""sayakpaul/unigram-tokenizer-wikitext"""					,   help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""					,   )
		parser.add_argument(
		    """--shard_size"""					,   type=UpperCamelCase__					,   default=1_000					,   help="""Number of entries to go in a single shard."""					,   )
		parser.add_argument("""--split"""					,   type=UpperCamelCase__					,   default="""train"""					,   choices=["""train""", """test""", """validation"""]    )
		parser.add_argument(
		    """--limit"""					,   default=UpperCamelCase__					,   type=UpperCamelCase__					,   help="""Limit the number of shards (used for debugging)."""					,   )
		parser.add_argument(
		    """--max_length"""					,   type=UpperCamelCase__					,   default=512					,   help="""Maximum sequence length. For training on TPUs, it helps to have a maximum"""
		    """ sequence length that is a multiple of 8."""					,   )
		parser.add_argument(
		    """--output_dir"""					,   default="""tf-tpu"""					,   type=UpperCamelCase__					,   help="""Output directory where the TFRecord shards will be saved. If the"""
		    """ path is appended with `gs://` ('gs://tf-tpu', for example) then the TFRecord"""
		    """ shards will be directly saved to a Google Cloud Storage bucket."""					,   )

		SCREAMING_SNAKE_CASE__			=      parser.parse_args()
		return args






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[Any]    ):
		def fn(UpperCamelCase__:							Any    ):
				return tokenizer(examples["""text"""]    )

		return fn






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							Any    ):
		SCREAMING_SNAKE_CASE__			=      []
		for i in range(len(tokenized_data["""input_ids"""]    )    ):
				SCREAMING_SNAKE_CASE__			=      {
				    """input_ids""": tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data["""input_ids"""][i]    )    ),
				    """attention_mask""": tf.train.Feature(
				        intaa_list=tf.train.IntaaList(value=tokenized_data["""attention_mask"""][i]    )    ),
				}
				SCREAMING_SNAKE_CASE__			=      tf.train.Features(feature=UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      tf.train.Example(features=UpperCamelCase__    )
				SCREAMING_SNAKE_CASE__			=      example.SerializeToString()
				records.append(UpperCamelCase__    )
		return records






def 				SCREAMING_SNAKE_CASE__							(       UpperCamelCase__:							List[str]    ):
		SCREAMING_SNAKE_CASE__			=      datasets.load_dataset(args.dataset_name					,   args.dataset_config					,   split=args.split    )

		if args.limit is not None:
				SCREAMING_SNAKE_CASE__			=      min(len(UpperCamelCase__    )					,   args.limit    )
				SCREAMING_SNAKE_CASE__			=      dataset.select(range(UpperCamelCase__    )    )
				print(f'''Limiting the dataset to {args.limit} entries.'''    )

		SCREAMING_SNAKE_CASE__			=      AutoTokenizer.from_pretrained(args.tokenizer_name_or_path    )

		# Handle output directory creation.
		# For serializing into a Google Cloud Storage Bucket, one needs to first
		# create a bucket.
		if "gs" not in args.output_dir:
				if not os.path.exists(args.output_dir    ):
						os.makedirs(args.output_dir    )
				SCREAMING_SNAKE_CASE__			=      os.path.join(args.output_dir					,   args.split    )
				if not os.path.exists(UpperCamelCase__    ):
						os.makedirs(UpperCamelCase__    )
		else:
				SCREAMING_SNAKE_CASE__			=      os.path.join(args.output_dir					,   args.split    )

		# Tokenize the whole dataset at once.
		SCREAMING_SNAKE_CASE__			=      tokenize_function(UpperCamelCase__    )
		SCREAMING_SNAKE_CASE__			=      dataset.map(UpperCamelCase__					,   batched=UpperCamelCase__					,   num_proc=4					,   remove_columns=["""text"""]    )

		# We need to concatenate all our texts together, and then split the result
		# into chunks of a fixed size, which we will call block_size. To do this, we
		# will use the map method again, with the option batched=True. When we use batched=True,
		# the function we pass to map() will be passed multiple inputs at once, allowing us
		# to group them into more or fewer examples than we had in the input.
		# This allows us to create our new fixed-length samples. The advantage of this
		# method is that we don't lose a whole lot of content from the dataset compared to the
		# case where we simply tokenize with a pre-defined max_length.

		def group_texts(UpperCamelCase__:							int    ):
				# Concatenate all texts.
				SCREAMING_SNAKE_CASE__			=      {k: sum(examples[k]					,   []    ) for k in examples.keys()}
				SCREAMING_SNAKE_CASE__			=      len(concatenated_examples[list(examples.keys()    )[0]]    )
				# We drop the small remainder, though you could add padding instead if the model supports it
				# In this, as in all things, we advise you to follow your heart 🫀
				SCREAMING_SNAKE_CASE__			=      (total_length // args.max_length) * args.max_length
				# Split by chunks of max_len.
				SCREAMING_SNAKE_CASE__			=      {
				    k: [t[i : i + args.max_length] for i in range(0					,   UpperCamelCase__					,   args.max_length    )]
				    for k, t in concatenated_examples.items()
				}
				return result

		SCREAMING_SNAKE_CASE__			=      dataset_tokenized.map(UpperCamelCase__					,   batched=UpperCamelCase__					,   batch_size=1_000					,   num_proc=4    )

		SCREAMING_SNAKE_CASE__			=      0
		SCREAMING_SNAKE_CASE__			=      0
		for shard in range(0					,   len(UpperCamelCase__    )					,   args.shard_size    ):
				SCREAMING_SNAKE_CASE__			=      grouped_dataset[shard : shard + args.shard_size]
				SCREAMING_SNAKE_CASE__			=      len(dataset_snapshot["""input_ids"""]    )
				SCREAMING_SNAKE_CASE__			=      os.path.join(UpperCamelCase__					,   f'''dataset-{shard_count}-{records_containing}.tfrecord'''    )
				SCREAMING_SNAKE_CASE__			=      get_serialized_examples(UpperCamelCase__    )

				with tf.io.TFRecordWriter(UpperCamelCase__    ) as out_file:
						for i in range(len(UpperCamelCase__    )    ):
								SCREAMING_SNAKE_CASE__			=      serialized_examples[i]
								out_file.write(UpperCamelCase__    )
						print("""Wrote file {} containing {} records""".format(UpperCamelCase__					,   UpperCamelCase__    )    )

				shard_count += 1
				total_records += records_containing

		with open(f'''split-{args.split}-records-count.txt'''					,   """w"""    ) as f:
				print(f'''Total {args.split} records: {total_records}'''					,   file=UpperCamelCase__    )


if __name__ == "__main__":
	_lowerCamelCase   				=			parse_args()
	main(args)