




'''simple docstring'''







import unittest

from transformers import XLMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
      import torch

      from transformers import (
          XLMForMultipleChoice,
          XLMForQuestionAnswering,
          XLMForQuestionAnsweringSimple,
          XLMForSequenceClassification,
          XLMForTokenClassification,
          XLMModel,
          XLMWithLMHeadModel,
      )
      from transformers.models.xlm.modeling_xlm import XLM_PRETRAINED_MODEL_ARCHIVE_LIST







class 					_a :





     '''simple docstring'''



     def __init__(				self							,__a							,__a=13							,__a=7							,__a=True							,__a=True							,__a=True							,__a=True							,__a=True							,__a=False							,__a=False							,__a=False							,__a=2							,__a=99							,__a=0							,__a=32							,__a=5							,__a=4							,__a=0.1							,__a=0.1							,__a=512							,__a=2							,__a=0.02							,__a=2							,__a=4							,__a="last"							,__a=True							,__a=None							,__a=0							,)      ->				List[Any]:
          snake_case			:      List[Any]        =			parent
          snake_case			:      List[str]        =			batch_size
          snake_case			:      Optional[int]        =			seq_length
          snake_case			:      List[str]        =			is_training
          snake_case			:      Tuple        =			use_input_lengths
          snake_case			:      Union[str, Any]        =			use_token_type_ids
          snake_case			:      Tuple        =			use_labels
          snake_case			:      List[str]        =			gelu_activation
          snake_case			:      Union[str, Any]        =			sinusoidal_embeddings
          snake_case			:      Optional[int]        =			causal
          snake_case			:      int        =			asm
          snake_case			:      List[str]        =			n_langs
          snake_case			:      Tuple        =			vocab_size
          snake_case			:      Any        =			n_special
          snake_case			:      Union[str, Any]        =			hidden_size
          snake_case			:      Optional[Any]        =			num_hidden_layers
          snake_case			:      Tuple        =			num_attention_heads
          snake_case			:      Any        =			hidden_dropout_prob
          snake_case			:      Dict        =			attention_probs_dropout_prob
          snake_case			:      Optional[Any]        =			max_position_embeddings
          snake_case			:      List[Any]        =			type_sequence_label_size
          snake_case			:      str        =			initializer_range
          snake_case			:      Dict        =			num_labels
          snake_case			:      Optional[int]        =			num_choices
          snake_case			:      Optional[Any]        =			summary_type
          snake_case			:      Dict        =			use_proj
          snake_case			:      Tuple        =			scope
          snake_case			:      List[str]        =			bos_token_id







     def 			snake_case_      (				self							)      ->				Optional[Any]:
          snake_case			:      int        =			ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size							)
          snake_case			:      List[Any]        =			random_attention_mask([self.batch_size, self.seq_length]							)

          snake_case			:      Union[str, 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			:      List[str]        =			None
          if self.use_token_type_ids:
               snake_case			:      str        =			ids_tensor([self.batch_size, self.seq_length]							,self.n_langs							)

          snake_case			:      Union[str, Any]        =			None
          snake_case			:      Tuple        =			None
          snake_case			:      List[Any]        =			None
          if self.use_labels:
               snake_case			:      Union[str, Any]        =			ids_tensor([self.batch_size]							,self.type_sequence_label_size							)
               snake_case			:      List[str]        =			ids_tensor([self.batch_size, self.seq_length]							,self.num_labels							)
               snake_case			:      Union[str, Any]        =			ids_tensor([self.batch_size]							,2							).float()
               snake_case			:      Union[str, Any]        =			ids_tensor([self.batch_size]							,self.num_choices							)

          snake_case			:      Optional[int]        =			self.get_config()

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







     def 			snake_case_      (				self							)      ->				List[str]:
          return 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							,num_labels=self.num_labels							,bos_token_id=self.bos_token_id							,)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,)      ->				int:
          snake_case			:      List[Any]        =			XLMModel(config=__a							)
          model.to(__a							)
          model.eval()
          snake_case			:      Tuple        =			model(__a							,lengths=__a							,langs=__a							)
          snake_case			:      str        =			model(__a							,langs=__a							)
          snake_case			:      List[str]        =			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							,__a							,)      ->				List[Any]:
          snake_case			:      Union[str, Any]        =			XLMWithLMHeadModel(__a							)
          model.to(__a							)
          model.eval()

          snake_case			:      List[str]        =			model(__a							,token_type_ids=__a							,labels=__a							)
          self.parent.assertEqual(result.loss.shape							,()							)
          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							,__a							,)      ->				Dict:
          snake_case			:      Optional[Any]        =			XLMForQuestionAnsweringSimple(__a							)
          model.to(__a							)
          model.eval()

          snake_case			:      int        =			model(__a							)

          snake_case			:      Tuple        =			model(__a							,start_positions=__a							,end_positions=__a							)
          snake_case			:      str        =			outputs
          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							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,)      ->				List[Any]:
          snake_case			:      Dict        =			XLMForQuestionAnswering(__a							)
          model.to(__a							)
          model.eval()

          snake_case			:      str        =			model(__a							)

          snake_case			:      Optional[int]        =			model(
              __a							,start_positions=__a							,end_positions=__a							,cls_index=__a							,is_impossible=__a							,p_mask=__a							,)

          snake_case			:      List[str]        =			model(
              __a							,start_positions=__a							,end_positions=__a							,cls_index=__a							,is_impossible=__a							,)

          ((snake_case) ,		)			:      Optional[int]        =			result_with_labels.to_tuple()

          snake_case			:      str        =			model(__a							,start_positions=__a							,end_positions=__a							)

          ((snake_case) ,		)			:      int        =			result_with_labels.to_tuple()

          self.parent.assertEqual(result_with_labels.loss.shape							,()							)
          self.parent.assertEqual(result.start_top_log_probs.shape							,(self.batch_size, model.config.start_n_top)							)
          self.parent.assertEqual(result.start_top_index.shape							,(self.batch_size, model.config.start_n_top)							)
          self.parent.assertEqual(
              result.end_top_log_probs.shape							,(self.batch_size, model.config.start_n_top * model.config.end_n_top)							)
          self.parent.assertEqual(
              result.end_top_index.shape							,(self.batch_size, model.config.start_n_top * model.config.end_n_top)							)
          self.parent.assertEqual(result.cls_logits.shape							,(self.batch_size,)							)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,)      ->				Tuple:
          snake_case			:      Optional[Any]        =			XLMForSequenceClassification(__a							)
          model.to(__a							)
          model.eval()

          snake_case			:      int        =			model(__a							)
          snake_case			:      Tuple        =			model(__a							,labels=__a							)
          self.parent.assertEqual(result.loss.shape							,()							)
          self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.type_sequence_label_size)							)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,__a							,)      ->				Optional[Any]:
          snake_case			:      Tuple        =			self.num_labels
          snake_case			:      Optional[int]        =			XLMForTokenClassification(__a							)
          model.to(__a							)
          model.eval()

          snake_case			:      str        =			model(__a							,attention_mask=__a							,labels=__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							,__a							,)      ->				Any:
          snake_case			:      Tuple        =			self.num_choices
          snake_case			:      List[Any]        =			XLMForMultipleChoice(config=__a							)
          model.to(__a							)
          model.eval()
          snake_case			:      List[Any]        =			input_ids.unsqueeze(1							).expand(-1							,self.num_choices							,-1							).contiguous()
          snake_case			:      Any        =			token_type_ids.unsqueeze(1							).expand(-1							,self.num_choices							,-1							).contiguous()
          snake_case			:      Union[str, Any]        =			input_mask.unsqueeze(1							).expand(-1							,self.num_choices							,-1							).contiguous()
          snake_case			:      Tuple        =			model(
              __a							,attention_mask=__a							,token_type_ids=__a							,labels=__a							,)
          self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_choices)							)







     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          snake_case			:      int        =			self.prepare_config_and_inputs()
          (
              (
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		(
              snake_case
          ) ,		
          )			:      List[str]        =			config_and_inputs
          snake_case			:      int        =			{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
          return config, inputs_dict







@require_torch
class 					_a (a__,		a__,		a__,		unittest.TestCase	):





     '''simple docstring'''



     lowerCAmelCase_  :     int          =					(
         (
             XLMModel,
             XLMWithLMHeadModel,
             XLMForQuestionAnswering,
             XLMForSequenceClassification,
             XLMForQuestionAnsweringSimple,
             XLMForTokenClassification,
             XLMForMultipleChoice,
         )
         if is_torch_available()
         else ()
     )
     lowerCAmelCase_  :     int          =					(
         (XLMWithLMHeadModel,) if is_torch_available() else ()
     )  # TODO (PVP): Check other models whether language generation is also applicable
     lowerCAmelCase_  :     Optional[Any]          =					(
         {
             """feature-extraction""": XLMModel,
             """fill-mask""": XLMWithLMHeadModel,
             """question-answering""": XLMForQuestionAnsweringSimple,
             """text-classification""": XLMForSequenceClassification,
             """text-generation""": XLMWithLMHeadModel,
             """token-classification""": XLMForTokenClassification,
             """zero-shot""": XLMForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							)      ->				int:
          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 			snake_case_      (				self							,__a							,__a							,__a=False							)      ->				Tuple:
          snake_case			:      Dict        =			super()._prepare_for_class(__a							,__a							,return_labels=__a							)

          if return_labels:
               if model_class.__name__ == "XLMForQuestionAnswering":
                    snake_case			:      Optional[int]        =			torch.zeros(
                        self.model_tester.batch_size							,dtype=torch.long							,device=__a							)
                    snake_case			:      Dict        =			torch.zeros(
                        self.model_tester.batch_size							,dtype=torch.long							,device=__a							)

          return inputs_dict







     def 			snake_case_      (				self							)      ->				List[Any]:
          snake_case			:      Any        =			XLMModelTester(self							)
          snake_case			:      Any        =			ConfigTester(self							,config_class=__a							,emb_dim=37							)







     def 			snake_case_      (				self							)      ->				int:
          self.config_tester.run_common_tests()







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      Union[str, Any]        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_model(*__a							)







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      List[Any]        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_lm_head(*__a							)







     def 			snake_case_      (				self							)      ->				str:
          snake_case			:      str        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_simple_qa(*__a							)







     def 			snake_case_      (				self							)      ->				Dict:
          snake_case			:      Any        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_qa(*__a							)







     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          snake_case			:      int        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_sequence_classif(*__a							)







     def 			snake_case_      (				self							)      ->				int:
          snake_case			:      List[str]        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_token_classif(*__a							)







     def 			snake_case_      (				self							)      ->				Any:
          snake_case			:      List[str]        =			self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_xlm_for_multiple_choice(*__a							)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							,__a=False							,__a=1							)      ->				int:
          self.assertIsInstance(__a							,__a							)
          self.assertListEqual(
              [isinstance(__a							,__a							) for iter_attentions in attentions]							,[True] * len(__a							)							)
          self.assertEqual(len(__a							)							,(max_length - min_length) * num_beam_groups							)

          for idx, iter_attentions in enumerate(__a							):
               # adds PAD dummy token
               snake_case			:      str        =			min_length + idx + 1
               snake_case			:      List[Any]        =			min_length + idx + 1

               snake_case			:      int        =			(
                   batch_size * num_beam_groups,
                   config.num_attention_heads,
                   tgt_len,
                   src_len,
               )
               # check attn size
               self.assertListEqual(
                   [layer_attention.shape for layer_attention in iter_attentions]							,[expected_shape] * len(__a							)							)







     def 			snake_case_      (				self							,__a							,__a							,__a							,__a							,__a							,__a=False							,__a=1							)      ->				str:
          self.assertIsInstance(__a							,__a							)
          self.assertListEqual(
              [isinstance(__a							,__a							) for iter_hidden_states in hidden_states]							,[True] * len(__a							)							,)
          self.assertEqual(len(__a							)							,(max_length - min_length) * num_beam_groups							)

          for idx, iter_hidden_states in enumerate(__a							):
               # adds PAD dummy token
               snake_case			:      List[str]        =			min_length + idx + 1
               snake_case			:      Optional[int]        =			(batch_size * num_beam_groups, seq_len, config.hidden_size)
               # check hidden size
               self.assertListEqual(
                   [layer_hidden_states.shape for layer_hidden_states in iter_hidden_states]							,[expected_shape] * len(__a							)							,)
          pass







     @slow
     def 			snake_case_      (				self							)      ->				Union[str, Any]:
          for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               snake_case			:      Any        =			XLMModel.from_pretrained(__a							)
               self.assertIsNotNone(__a							)







@require_torch
class 					_a (unittest.TestCase	):





     '''simple docstring'''



     @slow
     def 			snake_case_      (				self							)      ->				List[str]:
          snake_case			:      Tuple        =			XLMWithLMHeadModel.from_pretrained("""xlm-mlm-en-2048"""							)
          model.to(__a							)
          snake_case			:      Union[str, Any]        =			torch.tensor([[14, 447]]							,dtype=torch.long							,device=__a							)  # the president
          snake_case			:      Optional[Any]        =			[
              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			:      Dict        =			model.generate(__a							,do_sample=__a							)
          self.assertListEqual(output_ids[0].cpu().numpy().tolist()							,__a							)









'''simple docstring'''







import argparse
import logging
import pickle
from collections import Counter


logging.basicConfig(
    format="""%(asctime)s - %(levelname)s - %(name)s - %(message)s""", datefmt="""%m/%d/%Y %H:%M:%S""", level=logging.INFO
)
lowercase      :   Any      =       logging.getLogger(__name__)

if __name__ == "__main__":
      lowercase      :   List[Any]      =       argparse.ArgumentParser(
          description="""Token Counts for smoothing the masking probabilities in MLM (cf XLM/word2vec)"""
      )
      parser.add_argument(
          """--data_file""", type=str, default="""data/dump.bert-base-uncased.pickle""", help="""The binarized dataset."""
      )
      parser.add_argument(
          """--token_counts_dump""", type=str, default="""data/token_counts.bert-base-uncased.pickle""", help="""The dump file."""
      )
      parser.add_argument("""--vocab_size""", default=3_0522, type=int)
      lowercase      :   str      =       parser.parse_args()

      logger.info(F"""Loading data from {args.data_file}""")
      with open(args.data_file, """rb""") as fp:
            lowercase      :   int      =       pickle.load(fp)

      logger.info("""Counting occurrences for MLM.""")
      lowercase      :   List[Any]      =       Counter()
      for tk_ids in data:
            counter.update(tk_ids)
      lowercase      :   int      =       [0] * args.vocab_size
      for k, v in counter.items():
            lowercase      :   List[Any]      =       v

      logger.info(F"""Dump to {args.token_counts_dump}""")
      with open(args.token_counts_dump, """wb""") as handle:
            pickle.dump(counts, handle, protocol=pickle.HIGHEST_PROTOCOL)



