





"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
   import tensorflow as tf

   from transformers import (
       TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFXLMForMultipleChoice,
       TFXLMForQuestionAnsweringSimple,
       TFXLMForSequenceClassification,
       TFXLMForTokenClassification,
       TFXLMModel,
       TFXLMWithLMHeadModel,
       XLMConfig,
   )

class 			_UpperCAmelCase  :



      '''simple docstring'''

      def __init__(					self ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	List[Any]    						=  parent
        A_  :	Any    						=  1_3
        A_  :	Dict    						=  7
        A_  :	Dict    						=  True
        A_  :	Union[str, Any]    						=  True
        A_  :	str    						=  True
        A_  :	Union[str, Any]    						=  True
        A_  :	str    						=  True
        A_  :	Optional[Any]    						=  False
        A_  :	Optional[int]    						=  False
        A_  :	str    						=  False
        A_  :	str    						=  2
        A_  :	Tuple    						=  9_9
        A_  :	Optional[Any]    						=  0
        A_  :	Union[str, Any]    						=  3_2
        A_  :	List[Any]    						=  2
        A_  :	str    						=  4
        A_  :	int    						=  0.1
        A_  :	str    						=  0.1
        A_  :	Optional[Any]    						=  5_1_2
        A_  :	int    						=  1_6
        A_  :	Any    						=  2
        A_  :	Dict    						=  0.02
        A_  :	int    						=  3
        A_  :	str    						=  4
        A_  :	Union[str, Any]    						=  'last'
        A_  :	Dict    						=  True
        A_  :	Dict    						=  None
        A_  :	List[str]    						=  0

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Any    						=  ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size  )
        A_  :	Optional[int]    						=  random_attention_mask([self.batch_size, self.seq_length] ,				dtype=tf.floataa  )

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

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

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

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

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

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Tuple    						=  TFXLMModel(config=_A  )
        A_  :	Optional[int]    						=  {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
        A_  :	Any    						=  model(_A  )

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

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Any    						=  TFXLMWithLMHeadModel(_A  )

        A_  :	Dict    						=  {'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
        A_  :	str    						=  model(_A  )

        A_  :	str    						=  outputs

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

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Tuple    						=  TFXLMForQuestionAnsweringSimple(_A  )

        A_  :	Tuple    						=  {'input_ids': input_ids, 'lengths': input_lengths}

        A_  :	Dict    						=  model(_A  )

        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  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Optional[Any]    						=  TFXLMForSequenceClassification(_A  )

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

        A_  :	int    						=  model(_A  )

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

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Any    						=  self.num_labels
        A_  :	Optional[Any]    						=  TFXLMForTokenClassification(config=_A  )
        A_  :	List[Any]    						=  {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
        A_  :	Optional[int]    						=  model(_A  )
        self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)  )

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				):

        """simple docstring"""

        A_  :	Any    						=  self.num_choices
        A_  :	Dict    						=  TFXLMForMultipleChoice(config=_A  )
        A_  :	str    						=  tf.tile(tf.expand_dims(_A ,				1  ) ,				(1, self.num_choices, 1)  )
        A_  :	Union[str, Any]    						=  tf.tile(tf.expand_dims(_A ,				1  ) ,				(1, self.num_choices, 1)  )
        A_  :	Optional[Any]    						=  tf.tile(tf.expand_dims(_A ,				1  ) ,				(1, self.num_choices, 1)  )
        A_  :	Optional[int]    						=  {
            'input_ids': multiple_choice_inputs_ids,
            'attention_mask': multiple_choice_input_mask,
            'token_type_ids': multiple_choice_token_type_ids,
        }
        A_  :	Dict    						=  model(_A  )
        self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Dict    						=  self.prepare_config_and_inputs()
        (
            A_    
        )  :	List[str]    						=  config_and_inputs
        A_  :	Tuple    						=  {
            'input_ids': input_ids,
            'token_type_ids': token_type_ids,
            'langs': token_type_ids,
            'lengths': input_lengths,
        }
        return config, inputs_dict


@require_tf
class 			_UpperCAmelCase  (       snake_case__      ,       snake_case__      ,       unittest.TestCase		):



      '''simple docstring'''
      lowercase_				:			Optional[int]        =       (
          (
              TFXLMModel,
              TFXLMWithLMHeadModel,
              TFXLMForSequenceClassification,
              TFXLMForQuestionAnsweringSimple,
              TFXLMForTokenClassification,
              TFXLMForMultipleChoice,
          )
          if is_tf_available()
          else ()
      )
      lowercase_				:			str        =       (
          (TFXLMWithLMHeadModel,) if is_tf_available() else ()
      )  # TODO (PVP): Check other models whether language generation is also applicable
      lowercase_				:			str        =       (
          {
              '''feature-extraction''': TFXLMModel,
              '''fill-mask''': TFXLMWithLMHeadModel,
              '''question-answering''': TFXLMForQuestionAnsweringSimple,
              '''text-classification''': TFXLMForSequenceClassification,
              '''text-generation''': TFXLMWithLMHeadModel,
              '''token-classification''': TFXLMForTokenClassification,
              '''zero-shot''': TFXLMForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      lowercase_				:			Union[str, Any]        =       False
      lowercase_				:			Tuple        =       False

      def  lowerCamelCase_      (					self ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_ ,				snake_case_  ):

        """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  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Dict    						=  TFXLMModelTester(self  )
        A_  :	str    						=  ConfigTester(self ,				config_class=_A ,				emb_dim=3_7  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        self.config_tester.run_common_tests()

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

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

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Tuple    						=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_lm_head(*_A  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Optional[Any]    						=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_qa(*_A  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Dict    						=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_sequence_classif(*_A  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Union[str, Any]    						=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_token_classification(*_A  )

      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Optional[int]    						=  self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_xlm_for_multiple_choice(*_A  )

      @slow
      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
          A_  :	Tuple    						=  TFXLMModel.from_pretrained(_A  )
          self.assertIsNotNone(_A  )


@require_tf
class 			_UpperCAmelCase  (       unittest.TestCase		):



      '''simple docstring'''

      @slow
      def  lowerCamelCase_      (					self  ):

        """simple docstring"""

        A_  :	Tuple    						=  TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'  )
        A_  :	Dict    						=  tf.convert_to_tensor([[1_4, 4_4_7]] ,				dtype=tf.intaa  )  # the president
        A_  :	int    						=  [
            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_  :	str    						=  model.generate(_A ,				do_sample=_A  )
        self.assertListEqual(output_ids[0].numpy().tolist() ,				_A  )











"""simple docstring"""




from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if not inplace:
    A_  :	Optional[int]    						=  ma + ma
  else:
    ma += ma

  return ma





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  -1 * len(_UpperCAmelCase )
  A_  :	List[str]    						=  list(range(len(tensor.shape[:zero_index] ) ) )
  return tensor.permute(first_inds + [zero_index + i for i in inds] )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  return t.reshape(t.shape[:-no_dims] + (-1,) )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 1E-4 ):






  """simple docstring"""





  A_  :	Optional[Any]    						=  mask.expand(*value.shape )
  return torch.sum(mask * value			, dim=_UpperCAmelCase ) / (eps + torch.sum(_UpperCAmelCase			, dim=_UpperCAmelCase ))





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase = 2.3_125			, _UpperCAmelCase = 21.6_875			, _UpperCAmelCase = 64 ):






  """simple docstring"""





  A_  :	List[str]    						=  torch.linspace(_UpperCAmelCase			, _UpperCAmelCase			, no_bins - 1			, device=pts.device )
  A_  :	Optional[Any]    						=  torch.sqrt(torch.sum((pts.unsqueeze(-2 ) - pts.unsqueeze(-3 )) ** 2			, dim=-1 ) )
  return torch.bucketize(_UpperCAmelCase			, _UpperCAmelCase )





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Any    						=  dicts[0]
  A_  :	int    						=  {}
  for k, v in first.items():
    A_  :	Dict    						=  [d[k] for d in dicts]
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	List[Any]    						=  dict_multimap(_UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Union[str, Any]    						=  fn(_UpperCAmelCase )

  return new_dict





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Tuple    						=  v_bins.view(((1,) * len(x.shape )) + (len(_UpperCAmelCase ),) )
  A_  :	Any    						=  x[..., None] - reshaped_bins
  A_  :	List[str]    						=  torch.argmin(torch.abs(_UpperCAmelCase )			, dim=-1 )
  return nn.functional.one_hot(_UpperCAmelCase			, num_classes=len(_UpperCAmelCase ) ).float()





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase = 0			, _UpperCAmelCase = 0 ):






  """simple docstring"""





  A_  :	List[Union[slice, torch.Tensor]]    						=  []
  for i, s in enumerate(data.shape[:no_batch_dims] ):
    A_  :	List[Any]    						=  torch.arange(_UpperCAmelCase )
    A_  :	Dict    						=  r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape ) - i - 1))) )
    ranges.append(_UpperCAmelCase )

  A_  :	List[Union[slice, torch.Tensor]]    						=  [slice(_UpperCAmelCase ) for _ in range(len(data.shape ) - no_batch_dims )]
  A_  :	Optional[Any]    						=  inds
  ranges.extend(_UpperCAmelCase )
  # Matt note: Editing this to get around the behaviour of using a list as an array index changing
  # in recent Numpy versions
  return data[tuple(_UpperCAmelCase )]


lowerCamelCase_					:						Union[str, Any]      = TypeVar('T')





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  A_  :	Dict[Any, Union[dict, list, tuple, Any]]    						=  {}
  for k, v in dic.items():
    if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
      A_  :	int    						=  dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
    else:
      A_  :	Any    						=  tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )

  return new_dict





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





@overload
def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  ...





def 						UpperCAmelCase__		(					_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ):






  """simple docstring"""





  if isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return dict_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return [tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree]
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return tuple(tree_map(_UpperCAmelCase			, _UpperCAmelCase			, _UpperCAmelCase ) for x in tree )
  elif isinstance(_UpperCAmelCase			, _UpperCAmelCase ):
    return fn(_UpperCAmelCase )
  else:
    print(type(_UpperCAmelCase ) )
    raise ValueError('Not supported' )





lowerCamelCase_					:						Optional[int]      = partial(tree_map, leaf_type=torch.Tensor)