from __future__ import annotations

import unittest

from transformers import 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 (
				    TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
				    TFXLMForMultipleChoice,
				    TFXLMForQuestionAnsweringSimple,
				    TFXLMForSequenceClassification,
				    TFXLMForTokenClassification,
				    TFXLMModel,
				    TFXLMWithLMHeadModel,
				    XLMConfig,
				)



class    lowercase_    :



							def __init__(						self      ,				lowercase_      ,				):
													_snake_case						:     str				=    parent
													_snake_case						:     Tuple				=    13
													_snake_case						:     Dict				=    7
													_snake_case						:     Union[str, Any]				=    True
													_snake_case						:     Optional[Any]				=    True
													_snake_case						:     Tuple				=    True
													_snake_case						:     Dict				=    True
													_snake_case						:     List[Any]				=    True
													_snake_case						:     Union[str, Any]				=    False
													_snake_case						:     Any				=    False
													_snake_case						:     List[str]				=    False
													_snake_case						:     int				=    2
													_snake_case						:     Union[str, Any]				=    99
													_snake_case						:     Dict				=    0
													_snake_case						:     List[str]				=    32
													_snake_case						:     str				=    2
													_snake_case						:     int				=    4
													_snake_case						:     Dict				=    0.1
													_snake_case						:     int				=    0.1
													_snake_case						:     List[str]				=    512
													_snake_case						:     Union[str, Any]				=    16
													_snake_case						:     List[Any]				=    2
													_snake_case						:     int				=    0.02
													_snake_case						:     Union[str, Any]				=    3
													_snake_case						:     Union[str, Any]				=    4
													_snake_case						:     Optional[int]				=    'last'
													_snake_case						:     List[Any]				=    True
													_snake_case						:     Union[str, Any]				=    None
													_snake_case						:     int				=    0



							def        UpperCamelCase       (						self				):
													_snake_case						:     Tuple				=    ids_tensor([self.batch_size, self.seq_length]      ,				self.vocab_size				)
													_snake_case						:     Tuple				=    random_attention_mask([self.batch_size, self.seq_length]      ,				dtype=tf.floataa				)

													_snake_case						:     Optional[Any]				=    None
													if self.use_input_lengths:
																			_snake_case						:     Any				=    (
																			    ids_tensor([self.batch_size]      ,				vocab_size=2				) + self.seq_length - 2
																			)  # small variation of seq_length

													_snake_case						:     Optional[Any]				=    None
													if self.use_token_type_ids:
																			_snake_case						:     Tuple				=    ids_tensor([self.batch_size, self.seq_length]      ,				self.n_langs				)

													_snake_case						:     Any				=    None
													_snake_case						:     Union[str, Any]				=    None
													_snake_case						:     Optional[int]				=    None
													if self.use_labels:
																			_snake_case						:     Union[str, Any]				=    ids_tensor([self.batch_size]      ,				self.type_sequence_label_size				)
																			_snake_case						:     Any				=    ids_tensor([self.batch_size, self.seq_length]      ,				self.num_labels				)
																			_snake_case						:     Union[str, Any]				=    ids_tensor([self.batch_size]      ,				2      ,				dtype=tf.floataa				)
																			_snake_case						:     Any				=    ids_tensor([self.batch_size]      ,				self.num_choices				)

													_snake_case						:     int				=    XLMConfig(
													    vocab_size=self.vocab_size      ,				n_special=self.n_special      ,				emb_dim=self.hidden_size      ,				n_layers=self.num_hidden_layers      ,				n_heads=self.num_attention_heads      ,				dropout=self.hidden_dropout_prob      ,				attention_dropout=self.attention_probs_dropout_prob      ,				gelu_activation=self.gelu_activation      ,				sinusoidal_embeddings=self.sinusoidal_embeddings      ,				asm=self.asm      ,				causal=self.causal      ,				n_langs=self.n_langs      ,				max_position_embeddings=self.max_position_embeddings      ,				initializer_range=self.initializer_range      ,				summary_type=self.summary_type      ,				use_proj=self.use_proj      ,				bos_token_id=self.bos_token_id      ,				)

													return (
													    config,
													    input_ids,
													    token_type_ids,
													    input_lengths,
													    sequence_labels,
													    token_labels,
													    is_impossible_labels,
													    choice_labels,
													    input_mask,
													)



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     Tuple				=    TFXLMModel(config=snake_case_				)
													_snake_case						:     int				=    {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
													_snake_case						:     Tuple				=    model(snake_case_				)

													_snake_case						:     Any				=    [input_ids, input_mask]
													_snake_case						:     Optional[Any]				=    model(snake_case_				)
													self.parent.assertEqual(result.last_hidden_state.shape      ,				(self.batch_size, self.seq_length, self.hidden_size)				)



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     List[str]				=    TFXLMWithLMHeadModel(snake_case_				)

													_snake_case						:     Union[str, Any]				=    {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
													_snake_case						:     Optional[Any]				=    model(snake_case_				)

													_snake_case						:     Tuple				=    outputs

													self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.vocab_size)				)



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     List[str]				=    TFXLMForQuestionAnsweringSimple(snake_case_				)

													_snake_case						:     Optional[int]				=    {'input_ids': input_ids, 'lengths': input_lengths}

													_snake_case						:     Optional[int]				=    model(snake_case_				)

													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        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     List[str]				=    TFXLMForSequenceClassification(snake_case_				)

													_snake_case						:     Dict				=    {'input_ids': input_ids, 'lengths': input_lengths}

													_snake_case						:     Tuple				=    model(snake_case_				)

													self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.type_sequence_label_size)				)



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     int				=    self.num_labels
													_snake_case						:     List[str]				=    TFXLMForTokenClassification(config=snake_case_				)
													_snake_case						:     str				=    {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
													_snake_case						:     Optional[int]				=    model(snake_case_				)
													self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.seq_length, self.num_labels)				)



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				):
													_snake_case						:     List[Any]				=    self.num_choices
													_snake_case						:     Optional[Any]				=    TFXLMForMultipleChoice(config=snake_case_				)
													_snake_case						:     Tuple				=    tf.tile(tf.expand_dims(snake_case_      ,				1				)      ,				(1, self.num_choices, 1)				)
													_snake_case						:     Dict				=    tf.tile(tf.expand_dims(snake_case_      ,				1				)      ,				(1, self.num_choices, 1)				)
													_snake_case						:     str				=    tf.tile(tf.expand_dims(snake_case_      ,				1				)      ,				(1, self.num_choices, 1)				)
													_snake_case						:     Tuple				=    {
													    'input_ids': multiple_choice_inputs_ids,
													    'attention_mask': multiple_choice_input_mask,
													    'token_type_ids': multiple_choice_token_type_ids,
													}
													_snake_case						:     int				=    model(snake_case_				)
													self.parent.assertEqual(result.logits.shape      ,				(self.batch_size, self.num_choices)				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     Optional[int]				=    self.prepare_config_and_inputs()
													(
													    _snake_case			
													)						:     Optional[int]				=    config_and_inputs
													_snake_case						:     Any				=    {
													    'input_ids': input_ids,
													    'token_type_ids': token_type_ids,
													    'langs': token_type_ids,
													    'lengths': input_lengths,
													}
													return config, inputs_dict




@require_tf
class    lowercase_    (     lowerCamelCase_						, lowerCamelCase_						, unittest.TestCase		):
							_lowerCamelCase						=		(
							    (
							        TFXLMModel,
							        TFXLMWithLMHeadModel,
							        TFXLMForSequenceClassification,
							        TFXLMForQuestionAnsweringSimple,
							        TFXLMForTokenClassification,
							        TFXLMForMultipleChoice,
							    )
							    if is_tf_available()
							    else ()
							)
							_lowerCamelCase						=		(
							    (TFXLMWithLMHeadModel,) if is_tf_available() else ()
							)  # TODO (PVP): Check other models whether language generation is also applicable
							_lowerCamelCase						=		(
							    {
							        '''feature-extraction''': TFXLMModel,
							        '''fill-mask''': TFXLMWithLMHeadModel,
							        '''question-answering''': TFXLMForQuestionAnsweringSimple,
							        '''text-classification''': TFXLMForSequenceClassification,
							        '''text-generation''': TFXLMWithLMHeadModel,
							        '''token-classification''': TFXLMForTokenClassification,
							        '''zero-shot''': TFXLMForSequenceClassification,
							    }
							    if is_tf_available()
							    else {}
							)
							_lowerCamelCase						=		False
							_lowerCamelCase						=		False



							def        UpperCamelCase       (						self      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_      ,				lowercase_				):

													if (
													    pipeline_test_casse_name == "QAPipelineTests"
													    and tokenizer_name is not None
													    and not tokenizer_name.endswith("Fast"				)
													):
																			# `QAPipelineTests` fails for a few models when the slower tokenizer are used.
																			# (The slower tokenizers were never used for pipeline tests before the pipeline testing rework)
																			# TODO: check (and possibly fix) the `QAPipelineTests` with slower tokenizer
																			return True

													return False



							def        UpperCamelCase       (						self				):
													_snake_case						:     str				=    TFXLMModelTester(self				)
													_snake_case						:     int				=    ConfigTester(self      ,				config_class=snake_case_      ,				emb_dim=37				)



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



							def        UpperCamelCase       (						self				):
													_snake_case						:     Any				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_model(*snake_case_				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     Optional[int]				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_lm_head(*snake_case_				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     List[Any]				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_qa(*snake_case_				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     Tuple				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_sequence_classif(*snake_case_				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     Union[str, Any]				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_for_token_classification(*snake_case_				)



							def        UpperCamelCase       (						self				):
													_snake_case						:     Union[str, Any]				=    self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_xlm_for_multiple_choice(*snake_case_				)



							@slow
							def        UpperCamelCase       (						self				):

													for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																			_snake_case						:     int				=    TFXLMModel.from_pretrained(snake_case_				)
																			self.assertIsNotNone(snake_case_				)




@require_tf
class    lowercase_    (     unittest.TestCase		):



							@slow
							def        UpperCamelCase       (						self				):
													_snake_case						:     Any				=    TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"				)
													_snake_case						:     List[str]				=    tf.convert_to_tensor([[14, 447]]      ,				dtype=tf.intaa				)  # the president
													_snake_case						:     Optional[int]				=    [
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													    14,
													    447,
													]  # the president the president the president the president the president the president the president the president the president the president
													# TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
													_snake_case						:     Optional[int]				=    model.generate(snake_case_      ,				do_sample=snake_case_				)
													self.assertListEqual(output_ids[0].numpy().tolist()      ,				snake_case_				)
'''simple docstring'''


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


snake_case__     :			Optional[int]      			=  logging.getLogger(__name__)


def 					_lowerCamelCase    (      ):



     """simple docstring"""





     UpperCAmelCase_							:   List[Any]   =      argparse.ArgumentParser(description='Train a unigram tokenizer on 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(
         '--batch_size'	,     type=lowerCamelCase_	,     default=1000	,     help='Batch size during training.'	,     )
     parser.add_argument(
         '--vocab_size'	,     type=lowerCamelCase_	,     default=10048	,     help='Size of the desired vocabulary.'	,     )
     parser.add_argument(
         '--limit'	,     default=lowerCamelCase_	,     type=lowerCamelCase_	,     help='Limit the number of shards (used for debugging).'	,     )
     parser.add_argument(
         '--export_to_hub'	,     action='store_true'	,     )

     UpperCAmelCase_							:   str   =      parser.parse_args()
     return args


def 					_lowerCamelCase    (      lowerCamelCase_					:							Dict   ):



     """simple docstring"""





     UpperCAmelCase_							:   List[Any]   =      datasets.load_dataset(args.dataset_name	,     args.dataset_config	,     split='train'   )

     if args.limit is not None:
          UpperCAmelCase_							:   Union[str, Any]   =      min(len(lowerCamelCase_   )	,     args.limit   )
          UpperCAmelCase_							:   int   =      dataset.select(range(lowerCamelCase_   )   )
          logger.info(F'''Limiting the dataset to {args.limit} entries.'''   )

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

    # Prepare the tokenizer.
     UpperCAmelCase_							:   Optional[Any]   =      Tokenizer(Unigram()   )
     UpperCAmelCase_							:   Union[str, Any]   =      normalizers.Sequence([normalizers.Replace('``'	,     '"'   ), normalizers.Replace('\'\''	,     '"'   )]   )
     UpperCAmelCase_							:   Any   =      pre_tokenizers.Metaspace()

     # Prepare the trainer.
     UpperCAmelCase_							:   Tuple   =      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=lowerCamelCase_   )
     logger.info('Tokenizer training complete!'   )

     UpperCAmelCase_							:   List[str]   =      tokenizer.token_to_id('[CLS]'   )
     UpperCAmelCase_							:   Optional[int]   =      tokenizer.token_to_id('[SEP]'   )
     UpperCAmelCase_							:   Optional[Any]   =      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),
         ]	,     )
     UpperCAmelCase_							:   Any   =      decoders.Metaspace()

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


if __name__ == "__main__":
 snake_case__     :			Optional[Any]      			=  parse_args()
 main(args)






