"""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      _SCREAMING_SNAKE_CASE							:



   '''simple docstring'''



   def __init__(	self ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					parent
      __lowerCamelCase    		=					13
      __lowerCamelCase    		=					7
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					False
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					99
      __lowerCamelCase    		=					0
      __lowerCamelCase    		=					32
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					4
      __lowerCamelCase    		=					0.1
      __lowerCamelCase    		=					0.1
      __lowerCamelCase    		=					512
      __lowerCamelCase    		=					16
      __lowerCamelCase    		=					2
      __lowerCamelCase    		=					0.02
      __lowerCamelCase    		=					3
      __lowerCamelCase    		=					4
      __lowerCamelCase    		=					'last'
      __lowerCamelCase    		=					True
      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					0







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					ids_tensor([self.batch_size, self.seq_length] ,				self.vocab_size	)
      __lowerCamelCase    		=					random_attention_mask([self.batch_size, self.seq_length] ,				dtype=tf.floataa	)

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

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

      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					None
      __lowerCamelCase    		=					None
      if self.use_labels:
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				self.type_sequence_label_size	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size, self.seq_length] ,				self.num_labels	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				2 ,				dtype=tf.floataa	)
         __lowerCamelCase    		=					ids_tensor([self.batch_size] ,				self.num_choices	)

      __lowerCamelCase    		=					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    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMModel(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					[input_ids, input_mask]
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.last_hidden_state.shape ,				(self.batch_size, self.seq_length, self.hidden_size)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMWithLMHeadModel(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths, 'langs': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					outputs

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







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMForQuestionAnsweringSimple(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths}

      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

      self.parent.assertEqual(result.start_logits.shape ,				(self.batch_size, self.seq_length)	)
      self.parent.assertEqual(result.end_logits.shape ,				(self.batch_size, self.seq_length)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMForSequenceClassification(_SCREAMING_SNAKE_CASE	)

      __lowerCamelCase    		=					{'input_ids': input_ids, 'lengths': input_lengths}

      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)

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







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					self.num_labels
      __lowerCamelCase    		=					TFXLMForTokenClassification(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					{'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.seq_length, self.num_labels)	)







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				):


      '''simple docstring'''

      __lowerCamelCase    		=					self.num_choices
      __lowerCamelCase    		=					TFXLMForMultipleChoice(config=_SCREAMING_SNAKE_CASE	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					tf.tile(tf.expand_dims(_SCREAMING_SNAKE_CASE ,				1	) ,				(1, self.num_choices, 1)	)
      __lowerCamelCase    		=					{
          'input_ids': multiple_choice_inputs_ids,
          'attention_mask': multiple_choice_input_mask,
          'token_type_ids': multiple_choice_token_type_ids,
      }
      __lowerCamelCase    		=					model(_SCREAMING_SNAKE_CASE	)
      self.parent.assertEqual(result.logits.shape ,				(self.batch_size, self.num_choices)	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.prepare_config_and_inputs()
      (
          (
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,(
          __lowerCamelCase
      )		,
      )    		=					config_and_inputs
      __lowerCamelCase    		=					{
          'input_ids': input_ids,
          'token_type_ids': token_type_ids,
          'langs': token_type_ids,
          'lengths': input_lengths,
      }
      return config, inputs_dict






@require_tf
class      _SCREAMING_SNAKE_CASE							(			UpperCAmelCase    ,				UpperCAmelCase    ,				unittest.TestCase				):



   '''simple docstring'''



   __UpperCAmelCase	:							List[Any]    =    (
       (
           TFXLMModel,
           TFXLMWithLMHeadModel,
           TFXLMForSequenceClassification,
           TFXLMForQuestionAnsweringSimple,
           TFXLMForTokenClassification,
           TFXLMForMultipleChoice,
       )
       if is_tf_available()
       else ()
   )
   __UpperCAmelCase	:							List[str]    =    (
       (TFXLMWithLMHeadModel,) if is_tf_available() else ()
   )  # TODO (PVP): Check other models whether language generation is also applicable
   __UpperCAmelCase	:							List[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 {}
   )
   __UpperCAmelCase	:							Union[str, Any]    =    False
   __UpperCAmelCase	:							Optional[Any]    =    False







   def    _A				(	self ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_SNAKE_CASE ,				_SCREAMING_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    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMModelTester(self	)
      __lowerCamelCase    		=					ConfigTester(self ,				config_class=_SCREAMING_SNAKE_CASE ,				emb_dim=37	)







   def    _A				(	self	):


      '''simple docstring'''

      self.config_tester.run_common_tests()







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_model(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_lm_head(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_qa(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_sequence_classif(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_for_token_classification(*_SCREAMING_SNAKE_CASE	)







   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_xlm_for_multiple_choice(*_SCREAMING_SNAKE_CASE	)







   @slow
   def    _A				(	self	):


      '''simple docstring'''

      for model_name in TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         __lowerCamelCase    		=					TFXLMModel.from_pretrained(_SCREAMING_SNAKE_CASE	)
         self.assertIsNotNone(_SCREAMING_SNAKE_CASE	)






@require_tf
class      _SCREAMING_SNAKE_CASE							(			unittest.TestCase				):



   '''simple docstring'''



   @slow
   def    _A				(	self	):


      '''simple docstring'''

      __lowerCamelCase    		=					TFXLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'	)
      __lowerCamelCase    		=					tf.convert_to_tensor([[14, 447]] ,				dtype=tf.intaa	)  # the president
      __lowerCamelCase    		=					[
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
          14,
          447,
      ]  # the president the president the president the president the president the president the president the president the president the president
      # TODO(PVP): this and other input_ids I tried for generation give pretty bad results. Not sure why. Model might just not be made for auto-regressive inference
      __lowerCamelCase    		=					model.generate(_SCREAMING_SNAKE_CASE ,				do_sample=_SCREAMING_SNAKE_CASE	)
      self.assertListEqual(output_ids[0].numpy().tolist() ,				_SCREAMING_SNAKE_CASE	)












import json
import os
from collections.abc import Mapping
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from .imports import is_safetensors_available




def  a_		(							__snake_case     ,		__snake_case     ,		__snake_case     ,		__snake_case=None					)	->			str:

			'''simple docstring'''


			UpperCamelCase_     				= None
			# Check the string instead of the dtype to be compatible with versions of PyTorch that don't have bfloat16.
			if str(weight.dtype					) == "torch.bfloat16":
						# Need to reinterpret the underlined data as int16 since NumPy does not handle bfloat16s.
						UpperCamelCase_     				= weight.view(torch.intaa					)
						UpperCamelCase_     				= 'bfloat16'
			UpperCamelCase_     				= weight.cpu().numpy()
			UpperCamelCase_     				= os.path.join(__snake_case     ,		F'''{weight_name}.dat'''					)
			if index is not None:
						if dtype is None:
									UpperCamelCase_     				= str(array.dtype					)
						UpperCamelCase_     				= {'dtype': dtype, 'shape': list(array.shape					)}
			if array.ndim == 0:
						UpperCamelCase_     				= array[None]
			UpperCamelCase_     				= np.memmap(__snake_case     ,		dtype=array.dtype     ,		mode='w+'     ,		shape=array.shape					)
			UpperCamelCase_     				= array[:]
			file_array.flush()
			return index




def  a_		(							__snake_case     ,		__snake_case					)	->			Dict:

			'''simple docstring'''


			UpperCamelCase_     				= tuple(weight_info['shape']					)
			if shape == ():
						# NumPy memory-mapped arrays can't have 0 dims so it was saved as 1d tensor
						UpperCamelCase_     				= (1,)

			UpperCamelCase_     				= weight_info['dtype']
			if dtype == "bfloat16":
						# NumPy does not support bfloat16 so this was saved as a int16
						UpperCamelCase_     				= 'int16'

			UpperCamelCase_     				= np.memmap(__snake_case     ,		dtype=__snake_case     ,		shape=__snake_case     ,		mode='r'					)

			if len(weight_info['shape']					) == 0:
						UpperCamelCase_     				= weight[0]
			UpperCamelCase_     				= torch.tensor(__snake_case					)
			if weight_info["dtype"] == "bfloat16":
						UpperCamelCase_     				= weight.view(torch.bfloataa					)

			return weight




def  a_		(							__snake_case     ,		__snake_case					)	->			Optional[int]:

			'''simple docstring'''


			if index is None or len(__snake_case					) == 0:
						# Nothing to save
						return

			UpperCamelCase_     				= os.path.join(__snake_case     ,		'index.json'					)
			if os.path.isfile(__snake_case					):
						with open(__snake_case     ,		'r'     ,		encoding='utf-8'					) as f:
									UpperCamelCase_     				= json.load(__snake_case					)
			else:
						UpperCamelCase_     				= {}
			current_index.update(__snake_case					)

			with open(__snake_case     ,		'w'     ,		encoding='utf-8'					) as f:
						json.dump(__snake_case     ,		__snake_case     ,		indent=2					)




def  a_		(							__snake_case     ,		__snake_case					)	->			Tuple:

			'''simple docstring'''


			os.makedirs(__snake_case     ,		exist_ok=__snake_case					)
			UpperCamelCase_     				= {}
			for name, parameter in state_dict.items():
						UpperCamelCase_     				= offload_weight(__snake_case     ,		__snake_case     ,		__snake_case     ,		index=__snake_case					)

			# Update index
			save_offload_index(__snake_case     ,		__snake_case					)




class     A					(				lowerCamelCase_       ):







							def __init__(    self   :	int   ,     __UpperCAmelCase   :	Mapping   ,     __UpperCAmelCase   :	str				)		->					Tuple:

										"""simple docstring"""



										UpperCamelCase_     				= dataset
										UpperCamelCase_     				= prefix







							def __getitem__(    self   :	List[Any]   ,     __UpperCAmelCase   :	Dict				)		->					Tuple:

										"""simple docstring"""



										return self.dataset[f'''{self.prefix}{key}''']







							def __iter__(    self   :	int				)		->					List[str]:

										"""simple docstring"""



										return iter([key for key in self.dataset if key.startswith(self.prefix				)]				)







							def __len__(    self   :	Any				)		->					Optional[Any]:

										"""simple docstring"""



										return len(self.dataset				)







class     A					(				lowerCamelCase_       ):







							def __init__(    self   :	int   ,     __UpperCAmelCase   :	Dict[str, torch.Tensor] = None   ,     __UpperCAmelCase   :	Optional[Union[str, os.PathLike]] = None   ,     __UpperCAmelCase   :	Mapping = None   ,     __UpperCAmelCase   :	Optional[int]=None   ,     )		->					List[str]:

										"""simple docstring"""



										if state_dict is None and save_folder is None:
													raise ValueError('Need either a `state_dict` or a `save_folder` containing offloaded weights.'				)

										UpperCamelCase_     				= {} if state_dict is None else state_dict
										UpperCamelCase_     				= save_folder
										if index is None and save_folder is not None:
													with open(os.path.join(__UpperCAmelCase   ,     'index.json'				)				) as f:
																UpperCamelCase_     				= json.load(__UpperCAmelCase				)
										UpperCamelCase_     				= {} if index is None else index
										UpperCamelCase_     				= list(self.state_dict.keys()				)
										self.all_keys.extend([key for key in self.index if key not in self.all_keys]				)
										UpperCamelCase_     				= device







							def __getitem__(    self   :	Optional[Any]   ,     __UpperCAmelCase   :	str				)		->					Tuple:

										"""simple docstring"""



										if key in self.state_dict:
													return self.state_dict[key]
										UpperCamelCase_     				= self.index[key]
										if weight_info.get('safetensors_file'				) is not None:
													if not is_safetensors_available():
																raise ImportError('These offloaded weights require the use of safetensors: `pip install safetensors`.'				)

													from safetensors import safe_open

													UpperCamelCase_     				= 'cpu' if self.device is None else self.device
													with safe_open(weight_info['safetensors_file']   ,     framework='pt'   ,     device=__UpperCAmelCase				) as f:
																UpperCamelCase_     				= f.get_tensor(weight_info.get('weight_name'   ,     __UpperCAmelCase				)				)

													if "dtype" in weight_info:
																return tensor.to(getattr(__UpperCAmelCase   ,     weight_info['dtype']				)				)
													else:
																return tensor

										UpperCamelCase_     				= os.path.join(self.save_folder   ,     f'''{key}.dat'''				)
										return load_offloaded_weight(__UpperCAmelCase   ,     __UpperCAmelCase				)







							def __iter__(    self   :	Optional[int]				)		->					Union[str, Any]:

										"""simple docstring"""



										return iter(self.all_keys				)







							def __len__(    self   :	Any				)		->					Any:

										"""simple docstring"""



										return len(self.all_keys				)







def  a_		(							__snake_case     ,		__snake_case					)	->			Any:

			'''simple docstring'''


			UpperCamelCase_     				= {}
			for module_name in submodule_names:
						# We want to catch module_name parameter (module_name.xxx) or potentially module_name, but not any of the
						# submodules that could being like module_name (transformers.h.1 and transformers.h.10 for instance)
						result.update(
						    {
						        key: param
						        for key, param in state_dict.items()
						        if key == module_name or key.startswith(module_name + '.'					)
						    }					)
			return result
