

from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
					import tensorflow as tf

					from transformers.models.distilbert.modeling_tf_distilbert import (
					    TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
					    TFDistilBertForMaskedLM,
					    TFDistilBertForMultipleChoice,
					    TFDistilBertForQuestionAnswering,
					    TFDistilBertForSequenceClassification,
					    TFDistilBertForTokenClassification,
					    TFDistilBertModel,
					)





class       snake_case_  :







			'''simple docstring'''

			def __init__(	self					:							Optional[int]				,  _UpperCamelCase					:							int				,  ) ->Tuple:
									snake_case_  =      parent
									snake_case_  =      1_3
									snake_case_  =      7
									snake_case_  =      True
									snake_case_  =      True
									snake_case_  =      False
									snake_case_  =      True
									snake_case_  =      9_9
									snake_case_  =      3_2
									snake_case_  =      2
									snake_case_  =      4
									snake_case_  =      3_7
									snake_case_  =      '''gelu'''
									snake_case_  =      0.1
									snake_case_  =      0.1
									snake_case_  =      5_1_2
									snake_case_  =      1_6
									snake_case_  =      2
									snake_case_  =      0.02
									snake_case_  =      3
									snake_case_  =      4
									snake_case_  =      None







			def 							snake_case__(	self					:							Optional[Any]				) ->Any:
									snake_case_  =      ids_tensor([self.batch_size, self.seq_length]				,  self.vocab_size				)

									snake_case_  =      None
									if self.use_input_mask:
															snake_case_  =      random_attention_mask([self.batch_size, self.seq_length]				)

									snake_case_  =      None
									snake_case_  =      None
									snake_case_  =      None
									if self.use_labels:
															snake_case_  =      ids_tensor([self.batch_size]				,  self.type_sequence_label_size				)
															snake_case_  =      ids_tensor([self.batch_size, self.seq_length]				,  self.num_labels				)
															snake_case_  =      ids_tensor([self.batch_size]				,  self.num_choices				)

									snake_case_  =      DistilBertConfig(
									    vocab_size=self.vocab_size				,  dim=self.hidden_size				,  n_layers=self.num_hidden_layers				,  n_heads=self.num_attention_heads				,  hidden_dim=self.intermediate_size				,  hidden_act=self.hidden_act				,  dropout=self.hidden_dropout_prob				,  attention_dropout=self.attention_probs_dropout_prob				,  max_position_embeddings=self.max_position_embeddings				,  initializer_range=self.initializer_range				,  )

									return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







			def 							snake_case__(	self					:							List[str]				,  _UpperCamelCase					:							Optional[int]				,  _UpperCamelCase					:							Any				,  _UpperCamelCase					:							List[str]				,  _UpperCamelCase					:							List[Any]				,  _UpperCamelCase					:							int				,  _UpperCamelCase					:							Any				) ->List[str]:
									snake_case_  =      TFDistilBertModel(config=_UpperCamelCase				)
									snake_case_  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}

									snake_case_  =      model(_UpperCamelCase				)

									snake_case_  =      [input_ids, input_mask]

									snake_case_  =      model(_UpperCamelCase				)

									self.parent.assertEqual(result.last_hidden_state.shape				,  (self.batch_size, self.seq_length, self.hidden_size)				)







			def 							snake_case__(	self					:							Any				,  _UpperCamelCase					:							Tuple				,  _UpperCamelCase					:							Any				,  _UpperCamelCase					:							Dict				,  _UpperCamelCase					:							Optional[int]				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							str				) ->str:
									snake_case_  =      TFDistilBertForMaskedLM(config=_UpperCamelCase				)
									snake_case_  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
									snake_case_  =      model(_UpperCamelCase				)
									self.parent.assertEqual(result.logits.shape				,  (self.batch_size, self.seq_length, self.vocab_size)				)







			def 							snake_case__(	self					:							Union[str, Any]				,  _UpperCamelCase					:							Tuple				,  _UpperCamelCase					:							Tuple				,  _UpperCamelCase					:							Optional[Any]				,  _UpperCamelCase					:							int				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							Tuple				) ->Optional[int]:
									snake_case_  =      TFDistilBertForQuestionAnswering(config=_UpperCamelCase				)
									snake_case_  =      {
									    '''input_ids''': input_ids,
									    '''attention_mask''': input_mask,
									}
									snake_case_  =      model(_UpperCamelCase				)
									self.parent.assertEqual(result.start_logits.shape				,  (self.batch_size, self.seq_length)				)
									self.parent.assertEqual(result.end_logits.shape				,  (self.batch_size, self.seq_length)				)







			def 							snake_case__(	self					:							Dict				,  _UpperCamelCase					:							Tuple				,  _UpperCamelCase					:							List[str]				,  _UpperCamelCase					:							List[str]				,  _UpperCamelCase					:							Optional[Any]				,  _UpperCamelCase					:							Optional[Any]				,  _UpperCamelCase					:							Tuple				) ->str:
									snake_case_  =      self.num_labels
									snake_case_  =      TFDistilBertForSequenceClassification(_UpperCamelCase				)
									snake_case_  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
									snake_case_  =      model(_UpperCamelCase				)
									self.parent.assertEqual(result.logits.shape				,  (self.batch_size, self.num_labels)				)







			def 							snake_case__(	self					:							Union[str, Any]				,  _UpperCamelCase					:							int				,  _UpperCamelCase					:							Optional[int]				,  _UpperCamelCase					:							Tuple				,  _UpperCamelCase					:							str				,  _UpperCamelCase					:							Optional[int]				,  _UpperCamelCase					:							Optional[int]				) ->Optional[Any]:
									snake_case_  =      self.num_choices
									snake_case_  =      TFDistilBertForMultipleChoice(_UpperCamelCase				)
									snake_case_  =      tf.tile(tf.expand_dims(_UpperCamelCase				,  1				)				,  (1, self.num_choices, 1)				)
									snake_case_  =      tf.tile(tf.expand_dims(_UpperCamelCase				,  1				)				,  (1, self.num_choices, 1)				)
									snake_case_  =      {
									    '''input_ids''': multiple_choice_inputs_ids,
									    '''attention_mask''': multiple_choice_input_mask,
									}
									snake_case_  =      model(_UpperCamelCase				)
									self.parent.assertEqual(result.logits.shape				,  (self.batch_size, self.num_choices)				)







			def 							snake_case__(	self					:							Tuple				,  _UpperCamelCase					:							Optional[Any]				,  _UpperCamelCase					:							List[str]				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							Optional[int]				,  _UpperCamelCase					:							Union[str, Any]				,  _UpperCamelCase					:							Any				) ->Optional[Any]:
									snake_case_  =      self.num_labels
									snake_case_  =      TFDistilBertForTokenClassification(_UpperCamelCase				)
									snake_case_  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
									snake_case_  =      model(_UpperCamelCase				)
									self.parent.assertEqual(result.logits.shape				,  (self.batch_size, self.seq_length, self.num_labels)				)







			def 							snake_case__(	self					:							Union[str, Any]				) ->int:
									snake_case_  =      self.prepare_config_and_inputs()
									((snake_case_),       (snake_case_),       (snake_case_),       (snake_case_),       (snake_case_),       (snake_case_))  =      config_and_inputs
									snake_case_  =      {'''input_ids''': input_ids, '''attention_mask''': input_mask}
									return config, inputs_dict






@require_tf
class       snake_case_  (						__A      ,			__A      ,			unittest.TestCase   ):







			'''simple docstring'''

			SCREAMING_SNAKE_CASE	:  int		     =							(
			    (
			        TFDistilBertModel,
			        TFDistilBertForMaskedLM,
			        TFDistilBertForQuestionAnswering,
			        TFDistilBertForSequenceClassification,
			        TFDistilBertForTokenClassification,
			        TFDistilBertForMultipleChoice,
			    )
			    if is_tf_available()
			    else None
			)
			SCREAMING_SNAKE_CASE	:  List[Any]		     =							(
			    {
			        "feature-extraction": TFDistilBertModel,
			        "fill-mask": TFDistilBertForMaskedLM,
			        "question-answering": TFDistilBertForQuestionAnswering,
			        "text-classification": TFDistilBertForSequenceClassification,
			        "token-classification": TFDistilBertForTokenClassification,
			        "zero-shot": TFDistilBertForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			SCREAMING_SNAKE_CASE	:  Dict		     =							False
			SCREAMING_SNAKE_CASE	:  Optional[Any]		     =							False







			def 							snake_case__(	self					:							Optional[int]				) ->Dict:
									snake_case_  =      TFDistilBertModelTester(self				)
									snake_case_  =      ConfigTester(self				,  config_class=_UpperCamelCase				,  dim=3_7				)







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







			def 							snake_case__(	self					:							Tuple				) ->Union[str, Any]:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_model(*_UpperCamelCase				)







			def 							snake_case__(	self					:							str				) ->List[Any]:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_masked_lm(*_UpperCamelCase				)







			def 							snake_case__(	self					:							Optional[Any]				) ->List[Any]:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_question_answering(*_UpperCamelCase				)







			def 							snake_case__(	self					:							Tuple				) ->str:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_sequence_classification(*_UpperCamelCase				)







			def 							snake_case__(	self					:							Tuple				) ->Optional[Any]:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_multiple_choice(*_UpperCamelCase				)







			def 							snake_case__(	self					:							Optional[int]				) ->List[str]:
									snake_case_  =      self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_distilbert_for_token_classification(*_UpperCamelCase				)







			@slow
			def 							snake_case__(	self					:							List[str]				) ->Union[str, Any]:
									for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]				):
															snake_case_  =      TFDistilBertModel.from_pretrained(_UpperCamelCase				)
															self.assertIsNotNone(_UpperCamelCase				)






@require_tf
class       snake_case_  (						unittest.TestCase   ):







			'''simple docstring'''

			@slow
			def 							snake_case__(	self					:							Dict				) ->Optional[int]:
									snake_case_  =      TFDistilBertModel.from_pretrained('''distilbert-base-uncased'''				)
									snake_case_  =      tf.constant([[0, 1, 2, 3, 4, 5]]				)
									snake_case_  =      model(_UpperCamelCase				)[0]

									snake_case_  =      [1, 6, 7_6_8]
									self.assertEqual(output.shape				,  _UpperCamelCase				)

									snake_case_  =      tf.constant(
									    [
									        [
									            [0.19261885, -0.13732955, 0.4119799],
									            [0.22150156, -0.07422661, 0.39037204],
									            [0.22756018, -0.0896414, 0.3701467],
									        ]
									    ]				)
									tf.debugging.assert_near(output[:, :3, :3]				,  _UpperCamelCase				,  atol=1e-4				)



import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


UpperCAmelCase   :    Any     =		logging.getLogger(__name__)






def 	__lowerCamelCase						(							):





			'''simple docstring'''




			lowerCamelCase    		=  argparse.ArgumentParser(
			    description="""Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset."""      )
			parser.add_argument(
			    """--dataset_name"""  ,							type=lowerCamelCase__  ,							default="""wikitext"""  ,							help="""Name of the training. Explore datasets at: hf.co/datasets."""  ,							)
			parser.add_argument(
			    """--dataset_config"""  ,							type=lowerCamelCase__  ,							default="""wikitext-103-raw-v1"""  ,							help="""Configuration name of the dataset."""      )
			parser.add_argument(
			    """--tokenizer_name_or_path"""  ,							type=lowerCamelCase__  ,							default="""sayakpaul/unigram-tokenizer-wikitext"""  ,							help="""Tokenizer identifier. Can be a local filepath or a Hub identifier."""  ,							)
			parser.add_argument(
			    """--shard_size"""  ,							type=lowerCamelCase__  ,							default=1000  ,							help="""Number of entries to go in a single shard."""  ,							)
			parser.add_argument("""--split"""  ,							type=lowerCamelCase__  ,							default="""train"""  ,							choices=["""train""", """test""", """validation"""]      )
			parser.add_argument(
			    """--limit"""  ,							default=lowerCamelCase__  ,							type=lowerCamelCase__  ,							help="""Limit the number of shards (used for debugging)."""  ,							)
			parser.add_argument(
			    """--max_length"""  ,							type=lowerCamelCase__  ,							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=lowerCamelCase__  ,							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."""  ,							)

			lowerCamelCase    		=  parser.parse_args()
			return args






def 	__lowerCamelCase						(							lowerCamelCase__		:						Any      ):





			'''simple docstring'''




			def fn(lowerCamelCase__		:						Any      ):
						return tokenizer(examples["""text"""]      )

			return fn






def 	__lowerCamelCase						(							lowerCamelCase__		:						Optional[int]      ):





			'''simple docstring'''




			lowerCamelCase    		=  []
			for i in range(len(tokenized_data["""input_ids"""]      )      ):
						lowerCamelCase    		=  {
						    """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]      )      ),
						}
						lowerCamelCase    		=  tf.train.Features(feature=lowerCamelCase__      )
						lowerCamelCase    		=  tf.train.Example(features=lowerCamelCase__      )
						lowerCamelCase    		=  example.SerializeToString()
						records.append(lowerCamelCase__      )
			return records






def 	__lowerCamelCase						(							lowerCamelCase__		:						str      ):





			'''simple docstring'''




			lowerCamelCase    		=  datasets.load_dataset(args.dataset_name  ,							args.dataset_config  ,							split=args.split      )

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

			lowerCamelCase    		=  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      )
						lowerCamelCase    		=  os.path.join(args.output_dir  ,							args.split      )
						if not os.path.exists(lowerCamelCase__      ):
									os.makedirs(lowerCamelCase__      )
			else:
						lowerCamelCase    		=  os.path.join(args.output_dir  ,							args.split      )

			# Tokenize the whole dataset at once.
			lowerCamelCase    		=  tokenize_function(lowerCamelCase__      )
			lowerCamelCase    		=  dataset.map(lowerCamelCase__  ,							batched=lowerCamelCase__  ,							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(lowerCamelCase__		:						List[str]      ):
						# Concatenate all texts.
						lowerCamelCase    		=  {k: sum(examples[k]  ,							[]      ) for k in examples.keys()}
						lowerCamelCase    		=  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 🫀
						lowerCamelCase    		=  (total_length // args.max_length) * args.max_length
						# Split by chunks of max_len.
						lowerCamelCase    		=  {
						    k: [t[i : i + args.max_length] for i in range(0  ,							lowerCamelCase__  ,							args.max_length      )]
						    for k, t in concatenated_examples.items()
						}
						return result

			lowerCamelCase    		=  dataset_tokenized.map(lowerCamelCase__  ,							batched=lowerCamelCase__  ,							batch_size=1000  ,							num_proc=4      )

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

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

						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=lowerCamelCase__      )


if __name__ == "__main__":
		UpperCAmelCase   :    Optional[int]     =		parse_args()
		main(args)

