


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  SCREAMING_SNAKE_CASE__  :
      def __init__(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Tuple=1_3     ,							SCREAMING_SNAKE_CASE__				:       List[str]=7     ,							SCREAMING_SNAKE_CASE__				:       str=True     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=True     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=True     ,							SCREAMING_SNAKE_CASE__				:       Dict=True     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=True     ,							SCREAMING_SNAKE_CASE__				:       int=False     ,							SCREAMING_SNAKE_CASE__				:       str=False     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=False     ,							SCREAMING_SNAKE_CASE__				:       List[str]=2     ,							SCREAMING_SNAKE_CASE__				:       int=9_9     ,							SCREAMING_SNAKE_CASE__				:       str=0     ,							SCREAMING_SNAKE_CASE__				:       Dict=3_2     ,							SCREAMING_SNAKE_CASE__				:       str=5     ,							SCREAMING_SNAKE_CASE__				:       Any=4     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]=0.1     ,							SCREAMING_SNAKE_CASE__				:       List[str]=0.1     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=5_1_2     ,							SCREAMING_SNAKE_CASE__				:       Any=2     ,							SCREAMING_SNAKE_CASE__				:       List[Any]=0.02     ,							SCREAMING_SNAKE_CASE__				:       Any=2     ,							SCREAMING_SNAKE_CASE__				:       str=4     ,							SCREAMING_SNAKE_CASE__				:       Any="last"     ,							SCREAMING_SNAKE_CASE__				:       str=True     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]=None     ,							SCREAMING_SNAKE_CASE__				:       Tuple=0     ,							)		->			List[str]:
             a_					:		List[str]							=  parent
             a_					:		Any							=  batch_size
             a_					:		int							=  seq_length
             a_					:		Union[str, Any]							=  is_training
             a_					:		Dict							=  use_input_lengths
             a_					:		Union[str, Any]							=  use_token_type_ids
             a_					:		int							=  use_labels
             a_					:		Union[str, Any]							=  gelu_activation
             a_					:		List[str]							=  sinusoidal_embeddings
             a_					:		str							=  causal
             a_					:		Dict							=  asm
             a_					:		Dict							=  n_langs
             a_					:		List[Any]							=  vocab_size
             a_					:		Union[str, Any]							=  n_special
             a_					:		List[Any]							=  hidden_size
             a_					:		Optional[int]							=  num_hidden_layers
             a_					:		List[Any]							=  num_attention_heads
             a_					:		Optional[int]							=  hidden_dropout_prob
             a_					:		Tuple							=  attention_probs_dropout_prob
             a_					:		str							=  max_position_embeddings
             a_					:		Optional[Any]							=  type_sequence_label_size
             a_					:		Optional[Any]							=  initializer_range
             a_					:		List[Any]							=  num_labels
             a_					:		Union[str, Any]							=  num_choices
             a_					:		List[str]							=  summary_type
             a_					:		Tuple							=  use_proj
             a_					:		List[str]							=  scope
             a_					:		List[str]							=  bos_token_id
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Dict:
             a_					:		List[Any]							=  ids_tensor([self.batch_size, self.seq_length]     ,							self.vocab_size     )
             a_					:		List[str]							=  random_attention_mask([self.batch_size, self.seq_length]     )

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

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

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

             a_					:		Tuple							=  self.get_config()

             return (
                 config,
                 input_ids,
                 token_type_ids,
                 input_lengths,
                 sequence_labels,
                 token_labels,
                 is_impossible_labels,
                 choice_labels,
                 input_mask,
             )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			Optional[int]:
             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        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							)		->			List[str]:
             a_					:		List[Any]							=  XLMModel(config=__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()
             a_					:		List[Any]							=  model(__UpperCAmelCase     ,							lengths=__UpperCAmelCase     ,							langs=__UpperCAmelCase     )
             a_					:		Union[str, Any]							=  model(__UpperCAmelCase     ,							langs=__UpperCAmelCase     )
             a_					:		Any							=  model(__UpperCAmelCase     )
             self.parent.assertEqual(result.last_hidden_state.shape     ,							(self.batch_size, self.seq_length, self.hidden_size)     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							)		->			Any:
             a_					:		Union[str, Any]							=  XLMWithLMHeadModel(__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()

             a_					:		str							=  model(__UpperCAmelCase     ,							token_type_ids=__UpperCAmelCase     ,							labels=__UpperCAmelCase     )
             self.parent.assertEqual(result.loss.shape     ,							()     )
             self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.vocab_size)     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							)		->			Optional[Any]:
             a_					:		Union[str, Any]							=  XLMForQuestionAnsweringSimple(__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()

             a_					:		Dict							=  model(__UpperCAmelCase     )

             a_					:		Optional[int]							=  model(__UpperCAmelCase     ,							start_positions=__UpperCAmelCase     ,							end_positions=__UpperCAmelCase     )
             a_					:		Dict							=  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        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							)		->			List[Any]:
             a_					:		Any							=  XLMForQuestionAnswering(__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()

             a_					:		Union[str, Any]							=  model(__UpperCAmelCase     )

             a_					:		List[str]							=  model(
                 __UpperCAmelCase     ,							start_positions=__UpperCAmelCase     ,							end_positions=__UpperCAmelCase     ,							cls_index=__UpperCAmelCase     ,							is_impossible=__UpperCAmelCase     ,							p_mask=__UpperCAmelCase     ,							)

             a_					:		Union[str, Any]							=  model(
                 __UpperCAmelCase     ,							start_positions=__UpperCAmelCase     ,							end_positions=__UpperCAmelCase     ,							cls_index=__UpperCAmelCase     ,							is_impossible=__UpperCAmelCase     ,							)

             ((a_     )							,		     )					:		Optional[Any]							=  result_with_labels.to_tuple()

             a_					:		Dict							=  model(__UpperCAmelCase     ,							start_positions=__UpperCAmelCase     ,							end_positions=__UpperCAmelCase     )

             ((a_     )							,		     )					:		Optional[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        SCREAMING_SNAKE_CASE							(       self				:       Any     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							)		->			Dict:
             a_					:		List[Any]							=  XLMForSequenceClassification(__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()

             a_					:		Optional[Any]							=  model(__UpperCAmelCase     )
             a_					:		Any							=  model(__UpperCAmelCase     ,							labels=__UpperCAmelCase     )
             self.parent.assertEqual(result.loss.shape     ,							()     )
             self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.type_sequence_label_size)     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Any     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       str     ,							)		->			Optional[Any]:
             a_					:		int							=  self.num_labels
             a_					:		Any							=  XLMForTokenClassification(__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()

             a_					:		Any							=  model(__UpperCAmelCase     ,							attention_mask=__UpperCAmelCase     ,							labels=__UpperCAmelCase     )
             self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.seq_length, self.num_labels)     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							)		->			str:
             a_					:		Any							=  self.num_choices
             a_					:		Union[str, Any]							=  XLMForMultipleChoice(config=__UpperCAmelCase     )
             model.to(__UpperCAmelCase     )
             model.eval()
             a_					:		List[str]							=  input_ids.unsqueeze(1     ).expand(-1     ,							self.num_choices     ,							-1     ).contiguous()
             a_					:		Optional[Any]							=  token_type_ids.unsqueeze(1     ).expand(-1     ,							self.num_choices     ,							-1     ).contiguous()
             a_					:		Union[str, Any]							=  input_mask.unsqueeze(1     ).expand(-1     ,							self.num_choices     ,							-1     ).contiguous()
             a_					:		int							=  model(
                 __UpperCAmelCase     ,							attention_mask=__UpperCAmelCase     ,							token_type_ids=__UpperCAmelCase     ,							labels=__UpperCAmelCase     ,							)
             self.parent.assertEqual(result.logits.shape     ,							(self.batch_size, self.num_choices)     )

      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     )		->			Tuple:
             a_					:		List[Any]							=  self.prepare_config_and_inputs()
             (
                 (
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		(
                 a_     
             )							,		     
             )					:		Optional[Any]							=  config_and_inputs
             a_					:		Optional[int]							=  {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'lengths': input_lengths}
             return config, inputs_dict
@require_torch
class  SCREAMING_SNAKE_CASE__  (				lowercase__						,    lowercase__						,    lowercase__						,    unittest.TestCase			):
      snake_case__				:     str        =						(
          (
              XLMModel,
              XLMWithLMHeadModel,
              XLMForQuestionAnswering,
              XLMForSequenceClassification,
              XLMForQuestionAnsweringSimple,
              XLMForTokenClassification,
              XLMForMultipleChoice,
          )
          if is_torch_available()
          else ()
      )
      snake_case__				:     Any        =						(
          (XLMWithLMHeadModel,) if is_torch_available() else ()
      )  # TODO (PVP): Check other models whether language generation is also applicable
      snake_case__				:     str        =						(
          {
              '''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        SCREAMING_SNAKE_CASE							(       self				:       int     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       int     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Tuple     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     )		->			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        SCREAMING_SNAKE_CASE							(       self				:       Union[str, Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Any     ,							SCREAMING_SNAKE_CASE__				:       List[str]=False     )		->			int:
             a_					:		Dict							=  super()._prepare_for_class(__UpperCAmelCase     ,							__UpperCAmelCase     ,							return_labels=__UpperCAmelCase     )

             if return_labels:
                    if model_class.__name__ == "XLMForQuestionAnswering":
                           a_					:		Any							=  torch.zeros(
                               self.model_tester.batch_size     ,							dtype=torch.long     ,							device=__UpperCAmelCase     )
                           a_					:		Any							=  torch.zeros(
                               self.model_tester.batch_size     ,							dtype=torch.long     ,							device=__UpperCAmelCase     )

             return inputs_dict
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Any:
             a_					:		str							=  XLMModelTester(self     )
             a_					:		Optional[Any]							=  ConfigTester(self     ,							config_class=__UpperCAmelCase     ,							emb_dim=3_7     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Dict     )		->			Dict:
             self.config_tester.run_common_tests()
      def        SCREAMING_SNAKE_CASE							(       self				:       int     )		->			Optional[Any]:
             a_					:		Union[str, Any]							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_model(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			int:
             a_					:		int							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_lm_head(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Dict:
             a_					:		List[Any]							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_simple_qa(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Dict:
             a_					:		str							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_qa(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Tuple     )		->			Tuple:
             a_					:		Dict							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_sequence_classif(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[Any]     )		->			Optional[Any]:
             a_					:		Dict							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_token_classif(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Any     )		->			Union[str, Any]:
             a_					:		Tuple							=  self.model_tester.prepare_config_and_inputs()
             self.model_tester.create_and_check_xlm_for_multiple_choice(*__UpperCAmelCase     )
      def        SCREAMING_SNAKE_CASE							(       self				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Any     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       str     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       int=False     ,							SCREAMING_SNAKE_CASE__				:       int=1     )		->			Dict:
             self.assertIsInstance(__UpperCAmelCase     ,							__UpperCAmelCase     )
             self.assertListEqual(
                 [isinstance(__UpperCAmelCase     ,							__UpperCAmelCase     ) for iter_attentions in attentions]     ,							[True] * len(__UpperCAmelCase     )     )
             self.assertEqual(len(__UpperCAmelCase     )     ,							(max_length - min_length) * num_beam_groups     )

             for idx, iter_attentions in enumerate(__UpperCAmelCase     ):
                    # adds PAD dummy token
                    a_					:		int							=  min_length + idx + 1
                    a_					:		Optional[Any]							=  min_length + idx + 1

                    a_					:		Optional[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(__UpperCAmelCase     )     )
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[Any]     ,							SCREAMING_SNAKE_CASE__				:       List[Any]     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]     ,							SCREAMING_SNAKE_CASE__				:       List[str]     ,							SCREAMING_SNAKE_CASE__				:       Dict     ,							SCREAMING_SNAKE_CASE__				:       Optional[int]=False     ,							SCREAMING_SNAKE_CASE__				:       int=1     )		->			Dict:
             self.assertIsInstance(__UpperCAmelCase     ,							__UpperCAmelCase     )
             self.assertListEqual(
                 [isinstance(__UpperCAmelCase     ,							__UpperCAmelCase     ) for iter_hidden_states in hidden_states]     ,							[True] * len(__UpperCAmelCase     )     ,							)
             self.assertEqual(len(__UpperCAmelCase     )     ,							(max_length - min_length) * num_beam_groups     )

             for idx, iter_hidden_states in enumerate(__UpperCAmelCase     ):
                    # adds PAD dummy token
                    a_					:		int							=  min_length + idx + 1
                    a_					:		Dict							=  (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(__UpperCAmelCase     )     ,							)
             pass

      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       str     )		->			Tuple:


             for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                    a_					:		int							=  XLMModel.from_pretrained(__UpperCAmelCase     )
                    self.assertIsNotNone(__UpperCAmelCase     )
@require_torch
class  SCREAMING_SNAKE_CASE__  (				unittest.TestCase			):
      @slow
      def        SCREAMING_SNAKE_CASE							(       self				:       Optional[int]     )		->			int:
             a_					:		List[Any]							=  XLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'     )
             model.to(__UpperCAmelCase     )
             a_					:		Optional[int]							=  torch.tensor([[1_4, 4_4_7]]     ,							dtype=torch.long     ,							device=__UpperCAmelCase     )  # the president
             a_					:		Union[str, Any]							=  [
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
                 1_4,
                 4_4_7,
             ]  # 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
             a_					:		Optional[Any]							=  model.generate(__UpperCAmelCase     ,							do_sample=__UpperCAmelCase     )
             self.assertListEqual(output_ids[0].cpu().numpy().tolist()     ,							__UpperCAmelCase     )










import argparse
import json
from collections import OrderedDict

import torch
from huggingface_hub import cached_download, hf_hub_url

from transformers import AutoImageProcessor, CvtConfig, CvtForImageClassification
def 						_a    (      a :List[Any] )							->			Optional[int]:
 a		  =      []
 embed.append(
     (
         F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.weight""",
         F"""stage{idx}.patch_embed.proj.weight""",
     ) )
 embed.append(
     (
         F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.projection.bias""",
         F"""stage{idx}.patch_embed.proj.bias""",
     ) )
 embed.append(
     (
         F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.weight""",
         F"""stage{idx}.patch_embed.norm.weight""",
     ) )
 embed.append(
     (
         F"""cvt.encoder.stages.{idx}.embedding.convolution_embeddings.normalization.bias""",
         F"""stage{idx}.patch_embed.norm.bias""",
     ) )
 return embed
def 						_a    (      a :List[Any]	,							a :Optional[int] )							->			Dict:
 a		  =      []
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.convolution.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.conv.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_mean""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_mean""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.running_var""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.running_var""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_query.convolution_projection.normalization.num_batches_tracked""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_q.bn.num_batches_tracked""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.convolution.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.conv.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_mean""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_mean""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.running_var""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.running_var""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_key.convolution_projection.normalization.num_batches_tracked""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_k.bn.num_batches_tracked""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.convolution.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.conv.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_mean""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_mean""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.running_var""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.running_var""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.convolution_projection_value.convolution_projection.normalization.num_batches_tracked""",
         F"""stage{idx}.blocks.{cnt}.attn.conv_proj_v.bn.num_batches_tracked""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_q.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_query.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_q.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_k.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_key.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_k.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_v.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.attention.projection_value.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.proj_v.bias""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.weight""",
         F"""stage{idx}.blocks.{cnt}.attn.proj.weight""",
     ) )
 attention_weights.append(
     (
         F"""cvt.encoder.stages.{idx}.layers.{cnt}.attention.output.dense.bias""",
         F"""stage{idx}.blocks.{cnt}.attn.proj.bias""",
     ) )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.weight""", F"""stage{idx}.blocks.{cnt}.mlp.fc1.weight""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.intermediate.dense.bias""", F"""stage{idx}.blocks.{cnt}.mlp.fc1.bias""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.weight""", F"""stage{idx}.blocks.{cnt}.mlp.fc2.weight""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.output.dense.bias""", F"""stage{idx}.blocks.{cnt}.mlp.fc2.bias""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.weight""", F"""stage{idx}.blocks.{cnt}.norm1.weight""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_before.bias""", F"""stage{idx}.blocks.{cnt}.norm1.bias""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.weight""", F"""stage{idx}.blocks.{cnt}.norm2.weight""") )
 attention_weights.append(
     (F"""cvt.encoder.stages.{idx}.layers.{cnt}.layernorm_after.bias""", F"""stage{idx}.blocks.{cnt}.norm2.bias""") )
 return attention_weights
def 						_a    (      a :Any )							->			List[Any]:
 a		  =      []
 token.append((F"""cvt.encoder.stages.{idx}.cls_token""", '''stage2.cls_token''') )
 return token
def 						_a    (      )							->			Optional[int]:
 a		  =      []
 head.append(('''layernorm.weight''', '''norm.weight''') )
 head.append(('''layernorm.bias''', '''norm.bias''') )
 head.append(('''classifier.weight''', '''head.weight''') )
 head.append(('''classifier.bias''', '''head.bias''') )
 return head
def 						_a    (      a :Tuple	,							a :Optional[int]	,							a :List[Any]	,							a :Union[str, Any] )							->			Optional[int]:
 a		  =      '''imagenet-1k-id2label.json'''
 a		  =      1_000

 a		  =      '''huggingface/label-files'''
 a		  =      num_labels
 a		  =      json.load(open(cached_download(hf_hub_url(a	,							a	,							repo_type='''dataset''' ) )	,							'''r''' ) )
 a		  =      {int(a ): v for k, v in idalabel.items()}

 a		  =      idalabel
 a		  =      {v: k for k, v in idalabel.items()}

 a		  =      a		  =      CvtConfig(num_labels=a	,							idalabel=a	,							labelaid=a )

 # For depth size 13 (13 = 1+2+10)
 if cvt_model.rsplit('''/'''	,							1 )[-1][4:6] == "13":
  a		  =      [1, 2, 10]

 # For depth size 21 (21 = 1+4+16)
 elif cvt_model.rsplit('''/'''	,							1 )[-1][4:6] == "21":
  a		  =      [1, 4, 16]

 # For wide cvt (similar to wide-resnet) depth size 24 (w24 = 2 + 2 20)
 else:
  a		  =      [2, 2, 20]
  a		  =      [3, 12, 16]
  a		  =      [192, 768, 1_024]

 a		  =      CvtForImageClassification(a )
 a		  =      AutoImageProcessor.from_pretrained('''facebook/convnext-base-224-22k-1k''' )
 a		  =      image_size
 a		  =      torch.load(a	,							map_location=torch.device('''cpu''' ) )

 a		  =      OrderedDict()
 a		  =      []

 for idx in range(len(config.depth ) ):
  if config.cls_token[idx]:
   a		  =      list_of_state_dict + cls_token(a )
  a		  =      list_of_state_dict + embeddings(a )
  for cnt in range(config.depth[idx] ):
   a		  =      list_of_state_dict + attention(a	,							a )

 a		  =      list_of_state_dict + final()
 for gg in list_of_state_dict:
  print(a )
 for i in range(len(a ) ):
  a		  =      original_weights[list_of_state_dict[i][1]]

 model.load_state_dict(a )
 model.save_pretrained(a )
 image_processor.save_pretrained(a )


# Download the weights from zoo: https://1drv.ms/u/s!AhIXJn_J-blW9RzF3rMW7SsLHa8h?e=blQ0Al

if __name__ == "__main__":
     UpperCAmelCase__				   =		argparse.ArgumentParser()
     parser.add_argument(
         "--cvt_model",
         default="cvt-w24",
         type=str,
         help="Name of the cvt model you'd like to convert.",
     )
     parser.add_argument(
         "--image_size",
         default=384,
         type=int,
         help="Input Image Size",
     )
     parser.add_argument(
         "--cvt_file_name",
         default=R"cvtmodels\CvT-w24-384x384-IN-22k.pth",
         type=str,
         help="Input Image Size",
     )
     parser.add_argument(
         "--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model directory."
     )

     UpperCAmelCase__				   =		parser.parse_args()
     convert_cvt_checkpoint(args.cvt_model, args.image_size, args.cvt_file_name, args.pytorch_dump_folder_path)


