


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__:






      '''simple docstring'''







      def __init__(		self   ,			__SCREAMING_SNAKE_CASE   ,			)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			Optional[int] 						=parent
         UpperCamelCase__ :			Dict 						=13
         UpperCamelCase__ :			int 						=7
         UpperCamelCase__ :			Dict 						=True
         UpperCamelCase__ :			Tuple 						=True
         UpperCamelCase__ :			Tuple 						=True
         UpperCamelCase__ :			Any 						=True
         UpperCamelCase__ :			int 						=True
         UpperCamelCase__ :			Any 						=False
         UpperCamelCase__ :			Optional[Any] 						=False
         UpperCamelCase__ :			Dict 						=False
         UpperCamelCase__ :			Dict 						=2
         UpperCamelCase__ :			List[Any] 						=99
         UpperCamelCase__ :			Dict 						=0
         UpperCamelCase__ :			Any 						=32
         UpperCamelCase__ :			Optional[Any] 						=2
         UpperCamelCase__ :			Optional[Any] 						=4
         UpperCamelCase__ :			List[Any] 						=0.1
         UpperCamelCase__ :			Union[str, Any] 						=0.1
         UpperCamelCase__ :			Tuple 						=5_12
         UpperCamelCase__ :			List[Any] 						=16
         UpperCamelCase__ :			Union[str, Any] 						=2
         UpperCamelCase__ :			Any 						=0.02
         UpperCamelCase__ :			List[Any] 						=3
         UpperCamelCase__ :			Dict 						=4
         UpperCamelCase__ :			int 						="last"
         UpperCamelCase__ :			Any 						=True
         UpperCamelCase__ :			Union[str, Any] 						=None
         UpperCamelCase__ :			Dict 						=0




      def    UpperCAmelCase							(		self)					->						List[Any]:







         """simple docstring"""







         UpperCamelCase__ :			Dict 						=ids_tensor([self.batch_size, self.seq_length]   ,			self.vocab_size)
         UpperCamelCase__ :			List[str] 						=random_attention_mask([self.batch_size, self.seq_length]   ,			dtype=tf.floataa)

         UpperCamelCase__ :			Dict 						=None
         if self.use_input_lengths:
            UpperCamelCase__ :			List[str] 						=(
                ids_tensor([self.batch_size]   ,			vocab_size=2) + self.seq_length - 2
            )  # small variation of seq_length

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

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

         UpperCamelCase__ :			Dict 						=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   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						int:







         """simple docstring"""







         UpperCamelCase__ :			Union[str, Any] 						=TFXLMModel(config=__SCREAMING_SNAKE_CASE)
         UpperCamelCase__ :			Dict 						={"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
         UpperCamelCase__ :			Union[str, Any] 						=model(__SCREAMING_SNAKE_CASE)

         UpperCamelCase__ :			List[str] 						=[input_ids, input_mask]
         UpperCamelCase__ :			List[Any] 						=model(__SCREAMING_SNAKE_CASE)
         self.parent.assertEqual(result.last_hidden_state.shape   ,			(self.batch_size, self.seq_length, self.hidden_size))




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						Tuple:







         """simple docstring"""







         UpperCamelCase__ :			List[str] 						=TFXLMWithLMHeadModel(__SCREAMING_SNAKE_CASE)

         UpperCamelCase__ :			Any 						={"input_ids": input_ids, "lengths": input_lengths, "langs": token_type_ids}
         UpperCamelCase__ :			Optional[Any] 						=model(__SCREAMING_SNAKE_CASE)

         UpperCamelCase__ :			Optional[Any] 						=outputs

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




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			int 						=TFXLMForQuestionAnsweringSimple(__SCREAMING_SNAKE_CASE)

         UpperCamelCase__ :			Dict 						={"input_ids": input_ids, "lengths": input_lengths}

         UpperCamelCase__ :			Dict 						=model(__SCREAMING_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   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						int:







         """simple docstring"""







         UpperCamelCase__ :			List[str] 						=TFXLMForSequenceClassification(__SCREAMING_SNAKE_CASE)

         UpperCamelCase__ :			str 						={"input_ids": input_ids, "lengths": input_lengths}

         UpperCamelCase__ :			Optional[Any] 						=model(__SCREAMING_SNAKE_CASE)

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




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						List[Any]:







         """simple docstring"""







         UpperCamelCase__ :			Dict 						=self.num_labels
         UpperCamelCase__ :			Any 						=TFXLMForTokenClassification(config=__SCREAMING_SNAKE_CASE)
         UpperCamelCase__ :			Union[str, Any] 						={"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
         UpperCamelCase__ :			str 						=model(__SCREAMING_SNAKE_CASE)
         self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.seq_length, self.num_labels))




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			)					->						Tuple:







         """simple docstring"""







         UpperCamelCase__ :			List[str] 						=self.num_choices
         UpperCamelCase__ :			List[Any] 						=TFXLMForMultipleChoice(config=__SCREAMING_SNAKE_CASE)
         UpperCamelCase__ :			int 						=tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE   ,			1)   ,			(1, self.num_choices, 1))
         UpperCamelCase__ :			Dict 						=tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE   ,			1)   ,			(1, self.num_choices, 1))
         UpperCamelCase__ :			int 						=tf.tile(tf.expand_dims(__SCREAMING_SNAKE_CASE   ,			1)   ,			(1, self.num_choices, 1))
         UpperCamelCase__ :			Tuple 						={
             "input_ids": multiple_choice_inputs_ids,
             "attention_mask": multiple_choice_input_mask,
             "token_type_ids": multiple_choice_token_type_ids,
         }
         UpperCamelCase__ :			Any 						=model(__SCREAMING_SNAKE_CASE)
         self.parent.assertEqual(result.logits.shape   ,			(self.batch_size, self.num_choices))




      def    UpperCAmelCase							(		self)					->						Dict:







         """simple docstring"""







         UpperCamelCase__ :			Any 						=self.prepare_config_and_inputs()
         (
             (
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	(
             UpperCamelCase__
         )						,	
         ) :			str 						=config_and_inputs
         UpperCamelCase__ :			List[str] 						={
             "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__(						snake_case__			,							snake_case__			,							unittest.TestCase			):






      '''simple docstring'''







      snake_case__       =					(
          (
              TFXLMModel,
              TFXLMWithLMHeadModel,
              TFXLMForSequenceClassification,
              TFXLMForQuestionAnsweringSimple,
              TFXLMForTokenClassification,
              TFXLMForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      snake_case__       =					(
          (TFXLMWithLMHeadModel,) if is_tf_available() else ()
      )  # TODO (PVP): Check other models whether language generation is also applicable
      snake_case__       =					(
          {
              '''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 {}
      )
      snake_case__       =					False
      snake_case__       =					False




      def    UpperCAmelCase							(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE)					->						List[Any]:







         """simple docstring"""







         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)					->						str:







         """simple docstring"""







         UpperCamelCase__ :			List[Any] 						=TFXLMModelTester(self)
         UpperCamelCase__ :			int 						=ConfigTester(self   ,			config_class=__SCREAMING_SNAKE_CASE   ,			emb_dim=37)




      def    UpperCAmelCase							(		self)					->						Dict:







         """simple docstring"""







         self.config_tester.run_common_tests()




      def    UpperCAmelCase							(		self)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			List[Any] 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_model(*__SCREAMING_SNAKE_CASE)




      def    UpperCAmelCase							(		self)					->						Tuple:







         """simple docstring"""







         UpperCamelCase__ :			Optional[Any] 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_lm_head(*__SCREAMING_SNAKE_CASE)




      def    UpperCAmelCase							(		self)					->						Optional[int]:







         """simple docstring"""







         UpperCamelCase__ :			Tuple 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_qa(*__SCREAMING_SNAKE_CASE)




      def    UpperCAmelCase							(		self)					->						List[str]:







         """simple docstring"""







         UpperCamelCase__ :			Dict 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_sequence_classif(*__SCREAMING_SNAKE_CASE)




      def    UpperCAmelCase							(		self)					->						int:







         """simple docstring"""







         UpperCamelCase__ :			str 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_for_token_classification(*__SCREAMING_SNAKE_CASE)




      def    UpperCAmelCase							(		self)					->						Tuple:







         """simple docstring"""







         UpperCamelCase__ :			Union[str, Any] 						=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_xlm_for_multiple_choice(*__SCREAMING_SNAKE_CASE)




      @slow
      def    UpperCAmelCase							(		self)					->						List[str]:







         """simple docstring"""







         for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            UpperCamelCase__ :			Any 						=TFXLMModel.from_pretrained(__SCREAMING_SNAKE_CASE)
            self.assertIsNotNone(__SCREAMING_SNAKE_CASE)




@require_tf
class 							lowercase__(						unittest.TestCase			):






      '''simple docstring'''







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







         """simple docstring"""







         UpperCamelCase__ :			Dict 						=TFXLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048")
         UpperCamelCase__ :			Optional[Any] 						=tf.convert_to_tensor([[14, 4_47]]   ,			dtype=tf.intaa)  # the president
         UpperCamelCase__ :			Dict 						=[
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
             14,
             4_47,
         ]  # 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
         UpperCamelCase__ :			Optional[Any] 						=model.generate(__SCREAMING_SNAKE_CASE   ,			do_sample=__SCREAMING_SNAKE_CASE)
         self.assertListEqual(output_ids[0].numpy().tolist()   ,			__SCREAMING_SNAKE_CASE)





from collections import Counter
from functools import total_ordering

from data_structures.heap.heap import Heap



@total_ordering
class 							lowercase__:






      '''simple docstring'''







      def __init__(		self   ,			__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE)					->						None:







         """simple docstring"""







         UpperCamelCase__ :			Optional[int] 						=word
         UpperCamelCase__ :			Dict 						=count




      def __eq__(		self   ,			__SCREAMING_SNAKE_CASE)					->						bool:







         """simple docstring"""







         if not isinstance(__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE):
            return NotImplemented
         return self.count == other.count




      def __lt__(		self   ,			__SCREAMING_SNAKE_CASE)					->						bool:







         """simple docstring"""







         if not isinstance(__SCREAMING_SNAKE_CASE   ,			__SCREAMING_SNAKE_CASE):
            return NotImplemented
         return self.count < other.count




def 	_lowerCamelCase  (		A_  :	list[str]							,  A_  :	int      )			->			list[str]:

   '''simple docstring'''






   UpperCamelCase__ :			Heap[WordCount] 						=Heap()
   UpperCamelCase__ :			Any 						=Counter(A_      )
   heap.build_max_heap(
       [WordCount(A_							,  A_      ) for word, count in count_by_word.items()]      )
   return [heap.extract_max().word for _ in range(min(A_							,  len(A_      )      )      )]


if __name__ == "__main__":
     import doctest

     doctest.testmod()

