

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							:
  """simple docstring"""






  def __init__(      self			:   Optional[Any]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   str=13	,							lowerCAmelCase			:   Optional[Any]=7	,							lowerCAmelCase			:   str=True	,							lowerCAmelCase			:   Optional[Any]=True	,							lowerCAmelCase			:   Optional[Any]=True	,							lowerCAmelCase			:   List[Any]=True	,							lowerCAmelCase			:   Union[str, Any]=True	,							lowerCAmelCase			:   Dict=False	,							lowerCAmelCase			:   Any=False	,							lowerCAmelCase			:   Union[str, Any]=False	,							lowerCAmelCase			:   Optional[Any]=2	,							lowerCAmelCase			:   Tuple=99	,							lowerCAmelCase			:   Dict=0	,							lowerCAmelCase			:   str=32	,							lowerCAmelCase			:   Tuple=5	,							lowerCAmelCase			:   Any=4	,							lowerCAmelCase			:   List[str]=0.1	,							lowerCAmelCase			:   str=0.1	,							lowerCAmelCase			:   str=5_12	,							lowerCAmelCase			:   Optional[Any]=2	,							lowerCAmelCase			:   Tuple=0.02	,							lowerCAmelCase			:   Optional[int]=2	,							lowerCAmelCase			:   List[str]=4	,							lowerCAmelCase			:   Tuple="last"	,							lowerCAmelCase			:   Tuple=True	,							lowerCAmelCase			:   List[str]=None	,							lowerCAmelCase			:   List[str]=0	,							)							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						parent
     __lowerCAmelCase  :      str    =						batch_size
     __lowerCAmelCase  :      Tuple    =						seq_length
     __lowerCAmelCase  :      Tuple    =						is_training
     __lowerCAmelCase  :      Any    =						use_input_lengths
     __lowerCAmelCase  :      int    =						use_token_type_ids
     __lowerCAmelCase  :      List[str]    =						use_labels
     __lowerCAmelCase  :      int    =						gelu_activation
     __lowerCAmelCase  :      Dict    =						sinusoidal_embeddings
     __lowerCAmelCase  :      Dict    =						causal
     __lowerCAmelCase  :      Optional[int]    =						asm
     __lowerCAmelCase  :      str    =						n_langs
     __lowerCAmelCase  :      List[str]    =						vocab_size
     __lowerCAmelCase  :      Optional[Any]    =						n_special
     __lowerCAmelCase  :      Optional[Any]    =						hidden_size
     __lowerCAmelCase  :      Optional[Any]    =						num_hidden_layers
     __lowerCAmelCase  :      Tuple    =						num_attention_heads
     __lowerCAmelCase  :      Optional[Any]    =						hidden_dropout_prob
     __lowerCAmelCase  :      Any    =						attention_probs_dropout_prob
     __lowerCAmelCase  :      List[str]    =						max_position_embeddings
     __lowerCAmelCase  :      Optional[Any]    =						type_sequence_label_size
     __lowerCAmelCase  :      int    =						initializer_range
     __lowerCAmelCase  :      List[Any]    =						num_labels
     __lowerCAmelCase  :      Optional[Any]    =						num_choices
     __lowerCAmelCase  :      Any    =						summary_type
     __lowerCAmelCase  :      Optional[int]    =						use_proj
     __lowerCAmelCase  :      Optional[int]    =						scope
     __lowerCAmelCase  :      Optional[Any]    =						bos_token_id







  def 	SCREAMING_SNAKE_CASE   (      self			:   Union[str, Any]   )							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						ids_tensor([self.batch_size, self.seq_length]	,							self.vocab_size   )
     __lowerCAmelCase  :      Optional[Any]    =						random_attention_mask([self.batch_size, self.seq_length]   )

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

     __lowerCAmelCase  :      List[str]    =						None
     if self.use_token_type_ids:
        __lowerCAmelCase  :      str    =						ids_tensor([self.batch_size, self.seq_length]	,							self.n_langs   )

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

     __lowerCAmelCase  :      str    =						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			:   Union[str, Any]   )							->     Optional[int]:






     """simple docstring"""







     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			:   Any	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Optional[int]	,							)							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      int    =						XLMModel(config=lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()
     __lowerCAmelCase  :      Optional[Any]    =						model(lowerCAmelCase	,							lengths=lowerCAmelCase	,							langs=lowerCAmelCase   )
     __lowerCAmelCase  :      Tuple    =						model(lowerCAmelCase	,							langs=lowerCAmelCase   )
     __lowerCAmelCase  :      Any    =						model(lowerCAmelCase   )
     self.parent.assertEqual(result.last_hidden_state.shape	,							(self.batch_size, self.seq_length, self.hidden_size)   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   Any	,							)							->     List[str]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						XLMWithLMHeadModel(lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()

     __lowerCAmelCase  :      Tuple    =						model(lowerCAmelCase	,							token_type_ids=lowerCAmelCase	,							labels=lowerCAmelCase   )
     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			:   Optional[int]	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   List[Any]	,							)							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      Tuple    =						XLMForQuestionAnsweringSimple(lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()

     __lowerCAmelCase  :      Union[str, Any]    =						model(lowerCAmelCase   )

     __lowerCAmelCase  :      Optional[int]    =						model(lowerCAmelCase	,							start_positions=lowerCAmelCase	,							end_positions=lowerCAmelCase   )
     __lowerCAmelCase  :      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 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Optional[Any]	,							)							->     Dict:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						XLMForQuestionAnswering(lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()

     __lowerCAmelCase  :      Dict    =						model(lowerCAmelCase   )

     __lowerCAmelCase  :      List[str]    =						model(
         lowerCAmelCase	,							start_positions=lowerCAmelCase	,							end_positions=lowerCAmelCase	,							cls_index=lowerCAmelCase	,							is_impossible=lowerCAmelCase	,							p_mask=lowerCAmelCase	,							)

     __lowerCAmelCase  :      Optional[int]    =						model(
         lowerCAmelCase	,							start_positions=lowerCAmelCase	,							end_positions=lowerCAmelCase	,							cls_index=lowerCAmelCase	,							is_impossible=lowerCAmelCase	,							)

     (__lowerCAmelCase)  :      List[str]    =						result_with_labels.to_tuple()

     __lowerCAmelCase  :      Tuple    =						model(lowerCAmelCase	,							start_positions=lowerCAmelCase	,							end_positions=lowerCAmelCase   )

     (__lowerCAmelCase)  :      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			:   Optional[int]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   Optional[Any]	,							)							->     List[str]:






     """simple docstring"""







     __lowerCAmelCase  :      List[Any]    =						XLMForSequenceClassification(lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()

     __lowerCAmelCase  :      Dict    =						model(lowerCAmelCase   )
     __lowerCAmelCase  :      Optional[int]    =						model(lowerCAmelCase	,							labels=lowerCAmelCase   )
     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			:   Union[str, Any]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   int	,							)							->     int:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[Any]    =						self.num_labels
     __lowerCAmelCase  :      Tuple    =						XLMForTokenClassification(lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()

     __lowerCAmelCase  :      List[str]    =						model(lowerCAmelCase	,							attention_mask=lowerCAmelCase	,							labels=lowerCAmelCase   )
     self.parent.assertEqual(result.logits.shape	,							(self.batch_size, self.seq_length, self.num_labels)   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[str]	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   str	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Any	,							lowerCAmelCase			:   str	,							)							->     Tuple:






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.num_choices
     __lowerCAmelCase  :      Dict    =						XLMForMultipleChoice(config=lowerCAmelCase   )
     model.to(lowerCAmelCase   )
     model.eval()
     __lowerCAmelCase  :      List[Any]    =						input_ids.unsqueeze(1   ).expand(-1	,							self.num_choices	,							-1   ).contiguous()
     __lowerCAmelCase  :      Union[str, Any]    =						token_type_ids.unsqueeze(1   ).expand(-1	,							self.num_choices	,							-1   ).contiguous()
     __lowerCAmelCase  :      Union[str, Any]    =						input_mask.unsqueeze(1   ).expand(-1	,							self.num_choices	,							-1   ).contiguous()
     __lowerCAmelCase  :      int    =						model(
         lowerCAmelCase	,							attention_mask=lowerCAmelCase	,							token_type_ids=lowerCAmelCase	,							labels=lowerCAmelCase	,							)
     self.parent.assertEqual(result.logits.shape	,							(self.batch_size, self.num_choices)   )







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






     """simple docstring"""







     __lowerCAmelCase  :      int    =						self.prepare_config_and_inputs()
     (
         __lowerCAmelCase
     )  :      Any    =						config_and_inputs
     __lowerCAmelCase  :      Optional[Any]    =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
     return config, inputs_dict






@require_torch
class SCREAMING_SNAKE_CASE							(  a_       ,							a_       ,							a_       ,							unittest.TestCase			):
  """simple docstring"""






  lowerCamelCase   :      int										=(
      (
          XLMModel,
          XLMWithLMHeadModel,
          XLMForQuestionAnswering,
          XLMForSequenceClassification,
          XLMForQuestionAnsweringSimple,
          XLMForTokenClassification,
          XLMForMultipleChoice,
      )
      if is_torch_available()
      else ()
  )
  lowerCamelCase   :      str										=(
      (XLMWithLMHeadModel,) if is_torch_available() else ()
  )  # TODO (PVP): Check other models whether language generation is also applicable
  lowerCamelCase   :      Optional[int]										=(
      {
          "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			:   Dict	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   Union[str, Any]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   int	,							lowerCAmelCase			:   str   )							->     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 	SCREAMING_SNAKE_CASE   (      self			:   int	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Optional[int]=False   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						super()._prepare_for_class(lowerCAmelCase	,							lowerCAmelCase	,							return_labels=lowerCAmelCase   )

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

     return inputs_dict







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Any    =						XLMModelTester(self   )
     __lowerCAmelCase  :      Tuple    =						ConfigTester(self	,							config_class=lowerCAmelCase	,							emb_dim=37   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   int   )							->     Any:






     """simple docstring"""







     self.config_tester.run_common_tests()







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]   )							->     str:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_model(*lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Any:






     """simple docstring"""







     __lowerCAmelCase  :      List[str]    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_lm_head(*lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Tuple   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_simple_qa(*lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   str   )							->     Optional[int]:






     """simple docstring"""







     __lowerCAmelCase  :      List[str]    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_qa(*lowerCAmelCase   )







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






     """simple docstring"""







     __lowerCAmelCase  :      Optional[int]    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_sequence_classif(*lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]   )							->     List[Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Any    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_token_classif(*lowerCAmelCase   )







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






     """simple docstring"""







     __lowerCAmelCase  :      Dict    =						self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_xlm_for_multiple_choice(*lowerCAmelCase   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[int]	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Optional[Any]	,							lowerCAmelCase			:   Dict	,							lowerCAmelCase			:   Union[str, Any]=False	,							lowerCAmelCase			:   List[Any]=1   )							->     List[str]:






     """simple docstring"""







     self.assertIsInstance(lowerCAmelCase	,							lowerCAmelCase   )
     self.assertListEqual(
         [isinstance(lowerCAmelCase	,							lowerCAmelCase   ) for iter_attentions in attentions]	,							[True] * len(lowerCAmelCase   )   )
     self.assertEqual(len(lowerCAmelCase   )	,							(max_length - min_length) * num_beam_groups   )

     for idx, iter_attentions in enumerate(lowerCAmelCase   ):
        # adds PAD dummy token
        __lowerCAmelCase  :      Optional[Any]    =						min_length + idx + 1
        __lowerCAmelCase  :      Dict    =						min_length + idx + 1

        __lowerCAmelCase  :      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(lowerCAmelCase   )   )







  def 	SCREAMING_SNAKE_CASE   (      self			:   List[Any]	,							lowerCAmelCase			:   List[Any]	,							lowerCAmelCase			:   Optional[int]	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   Tuple	,							lowerCAmelCase			:   List[str]	,							lowerCAmelCase			:   int=False	,							lowerCAmelCase			:   Optional[int]=1   )							->     Dict:






     """simple docstring"""







     self.assertIsInstance(lowerCAmelCase	,							lowerCAmelCase   )
     self.assertListEqual(
         [isinstance(lowerCAmelCase	,							lowerCAmelCase   ) for iter_hidden_states in hidden_states]	,							[True] * len(lowerCAmelCase   )	,							)
     self.assertEqual(len(lowerCAmelCase   )	,							(max_length - min_length) * num_beam_groups   )

     for idx, iter_hidden_states in enumerate(lowerCAmelCase   ):
        # adds PAD dummy token
        __lowerCAmelCase  :      Dict    =						min_length + idx + 1
        __lowerCAmelCase  :      List[str]    =						(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(lowerCAmelCase   )	,							)
     pass







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






     """simple docstring"""







     for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        __lowerCAmelCase  :      Tuple    =						XLMModel.from_pretrained(lowerCAmelCase   )
        self.assertIsNotNone(lowerCAmelCase   )






@require_torch
class SCREAMING_SNAKE_CASE							(  unittest.TestCase			):
  """simple docstring"""






  @slow
  def 	SCREAMING_SNAKE_CASE   (      self			:   Optional[Any]   )							->     Union[str, Any]:






     """simple docstring"""







     __lowerCAmelCase  :      Union[str, Any]    =						XLMWithLMHeadModel.from_pretrained("""xlm-mlm-en-2048"""   )
     model.to(lowerCAmelCase   )
     __lowerCAmelCase  :      Dict    =						torch.tensor([[14, 4_47]]	,							dtype=torch.long	,							device=lowerCAmelCase   )  # the president
     __lowerCAmelCase  :      int    =						[
         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
     __lowerCAmelCase  :      Optional[int]    =						model.generate(lowerCAmelCase	,							do_sample=lowerCAmelCase   )
     self.assertListEqual(output_ids[0].cpu().numpy().tolist()	,							lowerCAmelCase   )





import argparse
from pathlib import Path

import requests
import torch
from PIL import Image

from transformers import (
    RobertaTokenizer,
    TrOCRConfig,
    TrOCRForCausalLM,
    TrOCRProcessor,
    VisionEncoderDecoderModel,
    ViTConfig,
    ViTImageProcessor,
    ViTModel,
)
from transformers.utils import logging


logging.set_verbosity_info()
__UpperCAmelCase								=			logging.get_logger(__name__)





def      snake_case_				(__A					:						Optional[int] ,       __A					:						Any					)      ->     Any:
   __lowerCAmelCase  :      Union[str, Any]    =						[]
   for i in range(encoder_config.num_hidden_layers					):
      # encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.norm1.weight''', f'''encoder.encoder.layer.{i}.layernorm_before.weight''')					)
      rename_keys.append((f'''encoder.deit.blocks.{i}.norm1.bias''', f'''encoder.encoder.layer.{i}.layernorm_before.bias''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.attn.proj.weight''', f'''encoder.encoder.layer.{i}.attention.output.dense.weight''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.attn.proj.bias''', f'''encoder.encoder.layer.{i}.attention.output.dense.bias''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.norm2.weight''', f'''encoder.encoder.layer.{i}.layernorm_after.weight''')					)
      rename_keys.append((f'''encoder.deit.blocks.{i}.norm2.bias''', f'''encoder.encoder.layer.{i}.layernorm_after.bias''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.mlp.fc1.weight''', f'''encoder.encoder.layer.{i}.intermediate.dense.weight''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.mlp.fc1.bias''', f'''encoder.encoder.layer.{i}.intermediate.dense.bias''')					)
      rename_keys.append(
          (f'''encoder.deit.blocks.{i}.mlp.fc2.weight''', f'''encoder.encoder.layer.{i}.output.dense.weight''')					)
      rename_keys.append((f'''encoder.deit.blocks.{i}.mlp.fc2.bias''', f'''encoder.encoder.layer.{i}.output.dense.bias''')					)

   # cls token, position embeddings and patch embeddings of encoder
   rename_keys.extend(
       [
           ("""encoder.deit.cls_token""", """encoder.embeddings.cls_token"""),
           ("""encoder.deit.pos_embed""", """encoder.embeddings.position_embeddings"""),
           ("""encoder.deit.patch_embed.proj.weight""", """encoder.embeddings.patch_embeddings.projection.weight"""),
           ("""encoder.deit.patch_embed.proj.bias""", """encoder.embeddings.patch_embeddings.projection.bias"""),
           ("""encoder.deit.norm.weight""", """encoder.layernorm.weight"""),
           ("""encoder.deit.norm.bias""", """encoder.layernorm.bias"""),
       ]					)

   return rename_keys





def      snake_case_				(__A					:						List[str] ,       __A					:						str					)      ->     Optional[Any]:
   for i in range(encoder_config.num_hidden_layers					):
      # queries, keys and values (only weights, no biases)
      __lowerCAmelCase  :      Optional[Any]    =						state_dict.pop(f'''encoder.deit.blocks.{i}.attn.qkv.weight'''					)

      __lowerCAmelCase  :      Tuple    =						in_proj_weight[
          : encoder_config.hidden_size, :
      ]
      __lowerCAmelCase  :      str    =						in_proj_weight[
          encoder_config.hidden_size : encoder_config.hidden_size * 2, :
      ]
      __lowerCAmelCase  :      str    =						in_proj_weight[
          -encoder_config.hidden_size :, :
      ]





def      snake_case_				(__A					:						Union[str, Any] ,       __A					:						str ,       __A					:						Optional[Any]					)      ->     Optional[Any]:
   __lowerCAmelCase  :      Any    =						dct.pop(__A					)
   __lowerCAmelCase  :      str    =						val





def      snake_case_				(__A					:						int					)      ->     Tuple:
   if "handwritten" in checkpoint_url:
      __lowerCAmelCase  :      Tuple    =						"""https://fki.tic.heia-fr.ch/static/img/a01-122-02-00.jpg"""  # industry
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-12.jpg" # have
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02-10.jpg" # let
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122-02.jpg"  #
      # url = "https://fki.tic.heia-fr.ch/static/img/a01-122.jpg"
   elif "printed" in checkpoint_url or "stage1" in checkpoint_url:
      __lowerCAmelCase  :      Optional[Any]    =						"""https://www.researchgate.net/profile/Dinh-Sang/publication/338099565/figure/fig8/AS:840413229350922@1577381536857/An-receipt-example-in-the-SROIE-2019-dataset_Q640.jpg"""
   __lowerCAmelCase  :      Dict    =						Image.open(requests.get(__A ,       stream=__A					).raw					).convert("""RGB"""					)
   return im





@torch.no_grad()
def      snake_case_				(__A					:						Any ,       __A					:						Union[str, Any]					)      ->     Optional[int]:
   __lowerCAmelCase  :      List[Any]    =						ViTConfig(image_size=3_8_4 ,       qkv_bias=__A					)
   __lowerCAmelCase  :      List[Any]    =						TrOCRConfig()

   # size of the architecture
   if "base" in checkpoint_url:
      __lowerCAmelCase  :      Union[str, Any]    =						7_6_8
   elif "large" in checkpoint_url:
      # use ViT-large encoder
      __lowerCAmelCase  :      Any    =						1_0_2_4
      __lowerCAmelCase  :      Any    =						4_0_9_6
      __lowerCAmelCase  :      Optional[int]    =						2_4
      __lowerCAmelCase  :      str    =						1_6
      __lowerCAmelCase  :      List[Any]    =						1_0_2_4
   else:
      raise ValueError("""Should either find 'base' or 'large' in checkpoint URL"""					)

   # the large-printed + stage1 checkpoints uses sinusoidal position embeddings, no layernorm afterwards
   if "large-printed" in checkpoint_url or "stage1" in checkpoint_url:
      __lowerCAmelCase  :      Tuple    =						False
      __lowerCAmelCase  :      Union[str, Any]    =						"""relu"""
      __lowerCAmelCase  :      List[Any]    =						1_0_2_4
      __lowerCAmelCase  :      Any    =						True
      __lowerCAmelCase  :      List[Any]    =						False
      __lowerCAmelCase  :      Dict    =						False

   # load HuggingFace model
   __lowerCAmelCase  :      Dict    =						ViTModel(__A ,       add_pooling_layer=__A					)
   __lowerCAmelCase  :      Union[str, Any]    =						TrOCRForCausalLM(__A					)
   __lowerCAmelCase  :      Any    =						VisionEncoderDecoderModel(encoder=__A ,       decoder=__A					)
   model.eval()

   # load state_dict of original model, rename some keys
   __lowerCAmelCase  :      Union[str, Any]    =						torch.hub.load_state_dict_from_url(__A ,       map_location="""cpu""" ,       check_hash=__A					)["""model"""]

   __lowerCAmelCase  :      Any    =						create_rename_keys(__A ,       __A					)
   for src, dest in rename_keys:
      rename_key(__A ,       __A ,       __A					)
   read_in_q_k_v(__A ,       __A					)

   # remove parameters we don't need
   del state_dict["encoder.deit.head.weight"]
   del state_dict["encoder.deit.head.bias"]
   del state_dict["decoder.version"]

   # add prefix to decoder keys
   for key, val in state_dict.copy().items():
      __lowerCAmelCase  :      Tuple    =						state_dict.pop(__A					)
      if key.startswith("""decoder"""					) and "output_projection" not in key:
         __lowerCAmelCase  :      str    =						val
      else:
         __lowerCAmelCase  :      Tuple    =						val

    # load state dict
   model.load_state_dict(__A					)

   # Check outputs on an image
   __lowerCAmelCase  :      List[Any]    =						ViTImageProcessor(size=encoder_config.image_size					)
   __lowerCAmelCase  :      List[str]    =						RobertaTokenizer.from_pretrained("""roberta-large"""					)
   __lowerCAmelCase  :      List[Any]    =						TrOCRProcessor(__A ,       __A					)

   __lowerCAmelCase  :      List[str]    =						processor(images=prepare_img(__A					) ,       return_tensors="""pt"""					).pixel_values

   # verify logits
   __lowerCAmelCase  :      List[str]    =						torch.tensor([[model.config.decoder.decoder_start_token_id]]					)
   __lowerCAmelCase  :      List[str]    =						model(pixel_values=__A ,       decoder_input_ids=__A					)
   __lowerCAmelCase  :      Optional[Any]    =						outputs.logits

   __lowerCAmelCase  :      Union[str, Any]    =						torch.Size([1, 1, 5_0_2_6_5]					)
   if "trocr-base-handwritten" in checkpoint_url:
      __lowerCAmelCase  :      Dict    =						torch.tensor(
          [-1.4502, -4.6683, -0.5347, -2.9291, 9.1435, -3.0571, 8.9764, 1.7560, 8.7358, -1.5311]					)
   elif "trocr-large-handwritten" in checkpoint_url:
      __lowerCAmelCase  :      List[Any]    =						torch.tensor(
          [-2.6437, -1.3129, -2.2596, -5.3455, 6.3539, 1.7604, 5.4991, 1.4702, 5.6113, 2.0170]					)
   elif "trocr-base-printed" in checkpoint_url:
      __lowerCAmelCase  :      Tuple    =						torch.tensor(
          [-5.6816, -5.8388, 1.1398, -6.9034, 6.8505, -2.4393, 1.2284, -1.0232, -1.9661, -3.9210]					)
   elif "trocr-large-printed" in checkpoint_url:
      __lowerCAmelCase  :      List[Any]    =						torch.tensor(
          [-6.0162, -7.0959, 4.4155, -5.1063, 7.0468, -3.1631, 2.6466, -0.3081, -0.8106, -1.7535]					)

   if "stage1" not in checkpoint_url:
      assert logits.shape == expected_shape, "Shape of logits not as expected"
      assert torch.allclose(logits[0, 0, :1_0] ,       __A ,       atol=1e-3					), "First elements of logits not as expected"

   Path(__A					).mkdir(exist_ok=__A					)
   print(f'''Saving model to {pytorch_dump_folder_path}'''					)
   model.save_pretrained(__A					)
   print(f'''Saving processor to {pytorch_dump_folder_path}'''					)
   processor.save_pretrained(__A					)


if __name__ == "__main__":
   __UpperCAmelCase								=			argparse.ArgumentParser()

   parser.add_argument(
       """--checkpoint_url""",
       default="""https://layoutlm.blob.core.windows.net/trocr/model_zoo/fairseq/trocr-base-handwritten.pt""",
       type=str,
       help="""URL to the original PyTorch checkpoint (.pth file).""",
   )
   parser.add_argument(
       """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the folder to output PyTorch model."""
   )
   __UpperCAmelCase								=			parser.parse_args()
   convert_tr_ocr_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


