

from __future__ import annotations

import unittest

from transformers import AutoTokenizer, MBartConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

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


if is_tf_available():
       import tensorflow as tf

       from transformers import TFAutoModelForSeqaSeqLM, TFMBartForConditionalGeneration, TFMBartModel
@require_tf
class      SCREAMING_SNAKE_CASE__  :





       '''simple docstring'''




       UpperCamelCase__			:  Tuple					     =							MBartConfig
       UpperCamelCase__			:  Any					     =							{}
       UpperCamelCase__			:  Optional[int]					     =							'''gelu'''

       def __init__(    self							:    Dict     ,		lowerCAmelCase__							:    Dict     ,		lowerCAmelCase__							:    int=13     ,		lowerCAmelCase__							:    List[Any]=7     ,		lowerCAmelCase__							:    Dict=True     ,		lowerCAmelCase__							:    List[str]=False     ,		lowerCAmelCase__							:    List[Any]=99     ,		lowerCAmelCase__							:    Dict=32     ,		lowerCAmelCase__							:    List[str]=2     ,		lowerCAmelCase__							:    Dict=4     ,		lowerCAmelCase__							:    List[Any]=37     ,		lowerCAmelCase__							:    Optional[int]=0.1     ,		lowerCAmelCase__							:    int=0.1     ,		lowerCAmelCase__							:    int=20     ,		lowerCAmelCase__							:    List[Any]=2     ,		lowerCAmelCase__							:    List[Any]=1     ,		lowerCAmelCase__							:    List[Any]=0     ,		)      ->     Dict:
             snake_case__											= parent
             snake_case__											= batch_size
             snake_case__											= seq_length
             snake_case__											= is_training
             snake_case__											= use_labels
             snake_case__											= vocab_size
             snake_case__											= hidden_size
             snake_case__											= num_hidden_layers
             snake_case__											= num_attention_heads
             snake_case__											= intermediate_size
             snake_case__											= hidden_dropout_prob
             snake_case__											= attention_probs_dropout_prob
             snake_case__											= max_position_embeddings
             snake_case__											= eos_token_id
             snake_case__											= pad_token_id
             snake_case__											= bos_token_id

       def UpperCAmelCase_    (    self							:    Optional[Any]				)      ->     Tuple:
             snake_case__											= ids_tensor([self.batch_size, self.seq_length - 1]     ,		self.vocab_size				)
             snake_case__											= tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size				)     ,		1				)
             snake_case__											= tf.concat([input_ids, eos_tensor]     ,		axis=1				)

             snake_case__											= ids_tensor([self.batch_size, self.seq_length]     ,		self.vocab_size				)

             snake_case__											= self.config_cls(
                 vocab_size=self.vocab_size     ,		d_model=self.hidden_size     ,		encoder_layers=self.num_hidden_layers     ,		decoder_layers=self.num_hidden_layers     ,		encoder_attention_heads=self.num_attention_heads     ,		decoder_attention_heads=self.num_attention_heads     ,		encoder_ffn_dim=self.intermediate_size     ,		decoder_ffn_dim=self.intermediate_size     ,		dropout=self.hidden_dropout_prob     ,		attention_dropout=self.attention_probs_dropout_prob     ,		max_position_embeddings=self.max_position_embeddings     ,		eos_token_ids=[2]     ,		bos_token_id=self.bos_token_id     ,		pad_token_id=self.pad_token_id     ,		decoder_start_token_id=self.pad_token_id     ,		**self.config_updates     ,		)
             snake_case__											= prepare_mbart_inputs_dict(UpperCAmelCase_     ,		UpperCAmelCase_     ,		UpperCAmelCase_				)
             return config, inputs_dict

       def UpperCAmelCase_    (    self							:    List[Any]     ,		lowerCAmelCase__							:    Any     ,		lowerCAmelCase__							:    Optional[Any]				)      ->     str:
             snake_case__											= TFMBartModel(config=UpperCAmelCase_				).get_decoder()
             snake_case__											= inputs_dict["""input_ids"""]

             snake_case__											= input_ids[:1, :]
             snake_case__											= inputs_dict["""attention_mask"""][:1, :]
             snake_case__											= inputs_dict["""head_mask"""]
             snake_case__											= 1

             # first forward pass
             snake_case__											= model(UpperCAmelCase_     ,		attention_mask=UpperCAmelCase_     ,		head_mask=UpperCAmelCase_     ,		use_cache=UpperCAmelCase_				)

             snake_case__						,      snake_case__											= outputs.to_tuple()
             snake_case__											= past_key_values[1]





def       _lowercase     (   __UpperCamelCase  :		Tuple ,							__UpperCamelCase  :		Union[str, Any] ,							__UpperCamelCase  :		int ,							__UpperCamelCase  :		Optional[int]=None ,							__UpperCamelCase  :		Tuple=None ,							__UpperCamelCase  :		Dict=None ,							__UpperCamelCase  :		Optional[int]=None ,							__UpperCamelCase  :		Optional[Any]=None ,							):
      if attention_mask is None:
            snake_case__											= tf.cast(tf.math.not_equal(UpperCAmelCase_ ,							config.pad_token_id			) ,							tf.inta			)
      if decoder_attention_mask is None:
            snake_case__											= tf.concat(
                [
                    tf.ones(decoder_input_ids[:, :1].shape ,							dtype=tf.inta			),
                    tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:] ,							config.pad_token_id			) ,							tf.inta			),
                ] ,							axis=-1 ,							)
      if head_mask is None:
            snake_case__											= tf.ones((config.encoder_layers, config.encoder_attention_heads)			)
      if decoder_head_mask is None:
            snake_case__											= tf.ones((config.decoder_layers, config.decoder_attention_heads)			)
      if cross_attn_head_mask is None:
            snake_case__											= tf.ones((config.decoder_layers, config.decoder_attention_heads)			)
      return {
          "input_ids": input_ids,
          "decoder_input_ids": decoder_input_ids,
          "attention_mask": attention_mask,
          "decoder_attention_mask": decoder_attention_mask,
          "head_mask": head_mask,
          "decoder_head_mask": decoder_head_mask,
          "cross_attn_head_mask": cross_attn_head_mask,
      }
@require_tf
class      SCREAMING_SNAKE_CASE__  (    _lowerCAmelCase			,						_lowerCAmelCase			,						unittest.TestCase		):





       '''simple docstring'''




       UpperCamelCase__			:  Optional[Any]					     =							(TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
       UpperCamelCase__			:  Optional[int]					     =							(TFMBartForConditionalGeneration,) if is_tf_available() else ()
       UpperCamelCase__			:  Any					     =							(
           {
               '''conversational''': TFMBartForConditionalGeneration,
               '''feature-extraction''': TFMBartModel,
               '''summarization''': TFMBartForConditionalGeneration,
               '''text2text-generation''': TFMBartForConditionalGeneration,
               '''translation''': TFMBartForConditionalGeneration,
           }
           if is_tf_available()
           else {}
       )
       UpperCamelCase__			:  Optional[Any]					     =							True
       UpperCamelCase__			:  str					     =							False
       UpperCamelCase__			:  int					     =							False

       def UpperCAmelCase_    (    self							:    Tuple     ,		lowerCAmelCase__							:    Any     ,		lowerCAmelCase__							:    List[Any]     ,		lowerCAmelCase__							:    List[Any]     ,		lowerCAmelCase__							:    Optional[Any]     ,		lowerCAmelCase__							:    List[Any]				)      ->     Dict:
             if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
                   # Exception encountered when calling layer '...'
                   return True

             return False

       def UpperCAmelCase_    (    self							:    int				)      ->     Optional[int]:
             snake_case__											= TFMBartModelTester(self				)
             snake_case__											= ConfigTester(self     ,		config_class=UpperCAmelCase_				)

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

       def UpperCAmelCase_    (    self							:    Optional[int]				)      ->     List[Any]:
             snake_case__											= self.model_tester.prepare_config_and_inputs_for_common()
             self.model_tester.check_decoder_model_past_large_inputs(*UpperCAmelCase_				)



@require_sentencepiece
@require_tokenizers
@require_tf
class      SCREAMING_SNAKE_CASE__  (    unittest.TestCase		):





       '''simple docstring'''




       UpperCamelCase__			:  Union[str, Any]					     =							[
           ''' UN Chief Says There Is No Military Solution in Syria''',
       ]
       UpperCamelCase__			:  Tuple					     =							[
           '''Şeful ONU declară că nu există o soluţie militară în Siria''',
       ]
       UpperCamelCase__			:  Optional[Any]					     =							'''facebook/mbart-large-en-ro'''

       @cached_property
       def UpperCAmelCase_    (    self							:    str				)      ->     Any:
             return AutoTokenizer.from_pretrained(self.model_name				)

       @cached_property
       def UpperCAmelCase_    (    self							:    Tuple				)      ->     Tuple:
             snake_case__											= TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name				)
             return model

       def UpperCAmelCase_    (    self							:    Optional[int]     ,		**lowerCAmelCase__							:    Optional[int]				)      ->     Optional[int]:
             snake_case__											= self.translate_src_text(**UpperCAmelCase_				)
             self.assertListEqual(self.expected_text     ,		UpperCAmelCase_				)

       def UpperCAmelCase_    (    self							:    List[str]     ,		**lowerCAmelCase__							:    int				)      ->     Dict:
             snake_case__											= self.tokenizer(self.src_text     ,		**UpperCAmelCase_     ,		return_tensors="""tf"""				)
             snake_case__											= self.model.generate(
                 model_inputs.input_ids     ,		attention_mask=model_inputs.attention_mask     ,		num_beams=2				)
             snake_case__											= self.tokenizer.batch_decode(UpperCAmelCase_     ,		skip_special_tokens=UpperCAmelCase_				)
             return generated_words

       @slow
       def UpperCAmelCase_    (    self							:    Optional[int]				)      ->     str:
             self._assert_generated_batch_equal_expected()








"""simple docstring"""


import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


SCREAMING_SNAKE_CASE			  =  logging.getLogger(__name__)



def 						lowerCamelCase__      (				)->				Tuple:







       """simple docstring"""
       UpperCamelCase       			=		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=10_00		,   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=5_12		,   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."		,   )

       UpperCamelCase       			=		parser.parse_args()
       return args



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Optional[int]:







       """simple docstring"""
       def fn(UpperCAmelCase_					):
              return tokenizer(examples["text"]					)

       return fn



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				Union[str, Any]:







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



def 						lowerCamelCase__      (				UpperCAmelCase_					)->				int:







       """simple docstring"""
       UpperCamelCase       			=		datasets.load_dataset(args.dataset_name		,   args.dataset_config		,   split=args.split					)

       if args.limit is not None:
              UpperCamelCase       			=		min(len(UpperCAmelCase_					)		,   args.limit					)
              UpperCamelCase       			=		dataset.select(range(UpperCAmelCase_					)					)
              print(F"Limiting the dataset to {args.limit} entries."					)

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

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

       UpperCamelCase       			=		dataset_tokenized.map(UpperCAmelCase_		,   batched=UpperCAmelCase_		,   batch_size=10_00		,   num_proc=4					)

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

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





