






import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


SCREAMING_SNAKE_CASE__				     =			logging.getLogger(__name__)

def        UpperCAmelCase__		(						):
			__a						:  Union[str, Any]									=     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=1_0_0_0   ,  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=5_1_2   ,  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.'   ,  )

			__a						:  int									=     parser.parse_args()
			return args

def        UpperCAmelCase__		(						lowerCamelCase_       :       Dict				):
			def fn(lowerCamelCase_       :       int				):
						return tokenizer(examples['text']				)

			return fn

def        UpperCAmelCase__		(						lowerCamelCase_       :       Dict				):
			__a						:  Dict									=     []
			for i in range(len(tokenized_data['input_ids']				)				):
						__a						:  Optional[int]									=     {
						    '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]				)				),
						}
						__a						:  Optional[Any]									=     tf.train.Features(feature=lowerCamelCase_				)
						__a						:  Dict									=     tf.train.Example(features=lowerCamelCase_				)
						__a						:  Any									=     example.SerializeToString()
						records.append(lowerCamelCase_				)
			return records

def        UpperCAmelCase__		(						lowerCamelCase_       :       Tuple				):
			__a						:  Optional[Any]									=     datasets.load_dataset(args.dataset_name   ,  args.dataset_config   ,  split=args.split				)

			if args.limit is not None:
						__a						:  Optional[int]									=     min(len(lowerCamelCase_				)   ,  args.limit				)
						__a						:  List[str]									=     dataset.select(range(lowerCamelCase_				)				)
						print(f'''Limiting the dataset to {args.limit} entries.'''				)

			__a						:  Any									=     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				)
						__a						:  Optional[Any]									=     os.path.join(args.output_dir   ,  args.split				)
						if not os.path.exists(lowerCamelCase_				):
									os.makedirs(lowerCamelCase_				)
			else:
						__a						:  Tuple									=     os.path.join(args.output_dir   ,  args.split				)

			# Tokenize the whole dataset at once.
			__a						:  Dict									=     tokenize_function(lowerCamelCase_				)
			__a						:  List[str]									=     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_       :       Tuple				):
						# Concatenate all texts.
						__a						:  Optional[int]									=     {k: sum(examples[k]   ,  []				) for k in examples.keys()}
						__a						:  Dict									=     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 🫀
						__a						:  int									=     (total_length // args.max_length) * args.max_length
						# Split by chunks of max_len.
						__a						:  int									=     {
						    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

			__a						:  Union[str, Any]									=     dataset_tokenized.map(lowerCamelCase_   ,  batched=lowerCamelCase_   ,  batch_size=1_0_0_0   ,  num_proc=4				)

			__a						:  Union[str, Any]									=     0
			__a						:  Union[str, Any]									=     0
			for shard in range(0   ,  len(lowerCamelCase_				)   ,  args.shard_size				):
						__a						:  List[Any]									=     grouped_dataset[shard : shard + args.shard_size]
						__a						:  Any									=     len(dataset_snapshot['input_ids']				)
						__a						:  str									=     os.path.join(lowerCamelCase_   ,  f'''dataset-{shard_count}-{records_containing}.tfrecord'''				)
						__a						:  int									=     get_serialized_examples(lowerCamelCase_				)

						with tf.io.TFRecordWriter(lowerCamelCase_				) as out_file:
									for i in range(len(lowerCamelCase_				)				):
												__a						:  Dict									=     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__":
							SCREAMING_SNAKE_CASE__				     =			parse_args()
							main(args)










from __future__ import annotations

import unittest

from transformers import DebertaVaConfig, 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 import (
       TFDebertaVaForMaskedLM,
       TFDebertaVaForQuestionAnswering,
       TFDebertaVaForSequenceClassification,
       TFDebertaVaForTokenClassification,
       TFDebertaVaModel,
   )



class      _UpperCAmelCase      :




     def __init__(					self			,    a__			,    a__=1_3			,    a__=7			,    a__=True			,    a__=True			,    a__=True			,    a__=True			,    a__=9_9			,    a__=3_2			,    a__=2			,    a__=4			,    a__=3_7			,    a__="gelu"			,    a__=0.1			,    a__=0.1			,    a__=5_1_2			,    a__=1_6			,    a__=2			,    a__=0.0_2			,    a__=False			,    a__=True			,    a__="None"			,    a__=3			,    a__=4			,    a__=None			,    ):
         A__					=       parent
         A__					=       batch_size
         A__					=       seq_length
         A__					=       is_training
         A__					=       use_input_mask
         A__					=       use_token_type_ids
         A__					=       use_labels
         A__					=       vocab_size
         A__					=       hidden_size
         A__					=       num_hidden_layers
         A__					=       num_attention_heads
         A__					=       intermediate_size
         A__					=       hidden_act
         A__					=       hidden_dropout_prob
         A__					=       attention_probs_dropout_prob
         A__					=       max_position_embeddings
         A__					=       type_vocab_size
         A__					=       type_sequence_label_size
         A__					=       initializer_range
         A__					=       num_labels
         A__					=       num_choices
         A__					=       relative_attention
         A__					=       position_biased_input
         A__					=       pos_att_type
         A__					=       scope




     def 		snake_case_		(					self):
         A__					=       ids_tensor([self.batch_size, self.seq_length]			,    self.vocab_size)

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

         A__					=       None
         if self.use_token_type_ids:
             A__					=       ids_tensor([self.batch_size, self.seq_length]			,    self.type_vocab_size)

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

         A__					=       DebertaVaConfig(
             vocab_size=self.vocab_size			,    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			,    max_position_embeddings=self.max_position_embeddings			,    type_vocab_size=self.type_vocab_size			,    relative_attention=self.relative_attention			,    position_biased_input=self.position_biased_input			,    initializer_range=self.initializer_range			,    return_dict=a__			,    )

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




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaModel(config=a__)
         A__					=       {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}

         A__					=       [input_ids, input_mask]
         A__					=       model(a__)

         A__					=       model(a__)

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




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaForMaskedLM(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.vocab_size))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       self.num_labels
         A__					=       TFDebertaVaForSequenceClassification(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.num_labels))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       self.num_labels
         A__					=       TFDebertaVaForTokenClassification(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }
         A__					=       model(a__)
         self.parent.assertEqual(result.logits.shape			,    (self.batch_size, self.seq_length, self.num_labels))




     def 		snake_case_		(					self			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__			,    a__):
         A__					=       TFDebertaVaForQuestionAnswering(config=a__)
         A__					=       {
             '''input_ids''': input_ids,
             '''attention_mask''': input_mask,
             '''token_type_ids''': token_type_ids,
         }

         A__					=       model(a__)
         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):
         A__					=       self.prepare_config_and_inputs()
         (
             (
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						(
             A__
         )			,						
         )					=       config_and_inputs
         A__					=       {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': input_mask}
         return config, inputs_dict







@require_tf
class      _UpperCAmelCase      (							A__      ,   A__      ,   unittest.TestCase  ):
     UpperCamelCase__		     =					(
         (
             TFDebertaVaModel,
             TFDebertaVaForMaskedLM,
             TFDebertaVaForQuestionAnswering,
             TFDebertaVaForSequenceClassification,
             TFDebertaVaForTokenClassification,
         )
         if is_tf_available()
         else ()
     )
     UpperCamelCase__		     =					(
         {
             '''feature-extraction''': TFDebertaVaModel,
             '''fill-mask''': TFDebertaVaForMaskedLM,
             '''question-answering''': TFDebertaVaForQuestionAnswering,
             '''text-classification''': TFDebertaVaForSequenceClassification,
             '''token-classification''': TFDebertaVaForTokenClassification,
             '''zero-shot''': TFDebertaVaForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )

     UpperCamelCase__		     =					False
     UpperCamelCase__		     =					False




     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModelTester(self)
         A__					=       ConfigTester(self			,    config_class=a__			,    hidden_size=3_7)




     def 		snake_case_		(					self):
         self.config_tester.run_common_tests()




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_masked_lm(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_question_answering(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_sequence_classification(*a__)




     def 		snake_case_		(					self):
         A__					=       self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_for_token_classification(*a__)




     @slow
     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge''')
         self.assertIsNotNone(a__)







@require_tf
class      _UpperCAmelCase      (							unittest.TestCase  ):




     @unittest.skip(reason='''Model not available yet''')
     def 		snake_case_		(					self):
         pass




     @slow
     def 		snake_case_		(					self):
         A__					=       TFDebertaVaModel.from_pretrained('''kamalkraj/deberta-v2-xlarge''')
         A__					=       tf.constant([[0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9, 4_6_0_7_8, 1_5_8_8, 2]])
         A__					=       tf.constant([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
         A__					=       model(a__			,    attention_mask=a__)[0]

         A__					=       tf.constant(
             [[[0.2_3_5_6, 0.1_9_4_8, 0.0_3_6_9], [-0.1_0_6_3, 0.3_5_8_6, -0.5_1_5_2], [-0.6_3_9_9, -0.0_2_5_9, -0.2_5_2_5]]])
         tf.debugging.assert_near(output[:, 1:4, 1:4]			,    a__			,    atol=1e-4)
