


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 							lowerCamelCase__						:






			"""simple docstring"""






			def __init__(				self   ,    snake_case   ,    snake_case=13   ,    snake_case=7   ,    snake_case=True   ,    snake_case=True   ,    snake_case=True   ,    snake_case=True   ,    snake_case=True   ,    snake_case=False   ,    snake_case=False   ,    snake_case=False   ,    snake_case=2   ,    snake_case=99   ,    snake_case=0   ,    snake_case=32   ,    snake_case=5   ,    snake_case=4   ,    snake_case=0.1   ,    snake_case=0.1   ,    snake_case=512   ,    snake_case=2   ,    snake_case=0.02   ,    snake_case=2   ,    snake_case=4   ,    snake_case="last"   ,    snake_case=True   ,    snake_case=None   ,    snake_case=0   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							parent
						UpperCamelCase__       =							batch_size
						UpperCamelCase__       =							seq_length
						UpperCamelCase__       =							is_training
						UpperCamelCase__       =							use_input_lengths
						UpperCamelCase__       =							use_token_type_ids
						UpperCamelCase__       =							use_labels
						UpperCamelCase__       =							gelu_activation
						UpperCamelCase__       =							sinusoidal_embeddings
						UpperCamelCase__       =							causal
						UpperCamelCase__       =							asm
						UpperCamelCase__       =							n_langs
						UpperCamelCase__       =							vocab_size
						UpperCamelCase__       =							n_special
						UpperCamelCase__       =							hidden_size
						UpperCamelCase__       =							num_hidden_layers
						UpperCamelCase__       =							num_attention_heads
						UpperCamelCase__       =							hidden_dropout_prob
						UpperCamelCase__       =							attention_probs_dropout_prob
						UpperCamelCase__       =							max_position_embeddings
						UpperCamelCase__       =							type_sequence_label_size
						UpperCamelCase__       =							initializer_range
						UpperCamelCase__       =							num_labels
						UpperCamelCase__       =							num_choices
						UpperCamelCase__       =							summary_type
						UpperCamelCase__       =							use_proj
						UpperCamelCase__       =							scope
						UpperCamelCase__       =							bos_token_id




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.vocab_size  )
						UpperCamelCase__       =							random_attention_mask([self.batch_size, self.seq_length]  )

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

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

						UpperCamelCase__       =							None
						UpperCamelCase__       =							None
						UpperCamelCase__       =							None
						if self.use_labels:
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    self.type_sequence_label_size  )
									UpperCamelCase__       =							ids_tensor([self.batch_size, self.seq_length]   ,    self.num_labels  )
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    2  ).float()
									UpperCamelCase__       =							ids_tensor([self.batch_size]   ,    self.num_choices  )

						UpperCamelCase__       =							self.get_config()

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




			def 					snake_case__					(				self  ):

						'''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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMModel(config=snake_case  )
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							model(snake_case   ,    lengths=snake_case   ,    langs=snake_case  )
						UpperCamelCase__       =							model(snake_case   ,    langs=snake_case  )
						UpperCamelCase__       =							model(snake_case  )
						self.parent.assertEqual(result.last_hidden_state.shape   ,    (self.batch_size, self.seq_length, self.hidden_size)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMWithLMHeadModel(snake_case  )
						model.to(snake_case  )
						model.eval()

						UpperCamelCase__       =							model(snake_case   ,    token_type_ids=snake_case   ,    labels=snake_case  )
						self.parent.assertEqual(result.loss.shape   ,    ()  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.vocab_size)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMForQuestionAnsweringSimple(snake_case  )
						model.to(snake_case  )
						model.eval()

						UpperCamelCase__       =							model(snake_case  )

						UpperCamelCase__       =							model(snake_case   ,    start_positions=snake_case   ,    end_positions=snake_case  )
						UpperCamelCase__       =							outputs
						self.parent.assertEqual(result.start_logits.shape   ,    (self.batch_size, self.seq_length)  )
						self.parent.assertEqual(result.end_logits.shape   ,    (self.batch_size, self.seq_length)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMForQuestionAnswering(snake_case  )
						model.to(snake_case  )
						model.eval()

						UpperCamelCase__       =							model(snake_case  )

						UpperCamelCase__       =							model(
						    snake_case   ,    start_positions=snake_case   ,    end_positions=snake_case   ,    cls_index=snake_case   ,    is_impossible=snake_case   ,    p_mask=snake_case   ,    )

						UpperCamelCase__       =							model(
						    snake_case   ,    start_positions=snake_case   ,    end_positions=snake_case   ,    cls_index=snake_case   ,    is_impossible=snake_case   ,    )

						(UpperCamelCase__				)       =							result_with_labels.to_tuple()

						UpperCamelCase__       =							model(snake_case   ,    start_positions=snake_case   ,    end_positions=snake_case  )

						(UpperCamelCase__				)       =							result_with_labels.to_tuple()

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




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMForSequenceClassification(snake_case  )
						model.to(snake_case  )
						model.eval()

						UpperCamelCase__       =							model(snake_case  )
						UpperCamelCase__       =							model(snake_case   ,    labels=snake_case  )
						self.parent.assertEqual(result.loss.shape   ,    ()  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.type_sequence_label_size)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							self.num_labels
						UpperCamelCase__       =							XLMForTokenClassification(snake_case  )
						model.to(snake_case  )
						model.eval()

						UpperCamelCase__       =							model(snake_case   ,    attention_mask=snake_case   ,    labels=snake_case  )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.seq_length, self.num_labels)  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    ):

						'''simple docstring'''


						UpperCamelCase__       =							self.num_choices
						UpperCamelCase__       =							XLMForMultipleChoice(config=snake_case  )
						model.to(snake_case  )
						model.eval()
						UpperCamelCase__       =							input_ids.unsqueeze(1  ).expand(-1   ,    self.num_choices   ,    -1  ).contiguous()
						UpperCamelCase__       =							token_type_ids.unsqueeze(1  ).expand(-1   ,    self.num_choices   ,    -1  ).contiguous()
						UpperCamelCase__       =							input_mask.unsqueeze(1  ).expand(-1   ,    self.num_choices   ,    -1  ).contiguous()
						UpperCamelCase__       =							model(
						    snake_case   ,    attention_mask=snake_case   ,    token_type_ids=snake_case   ,    labels=snake_case   ,    )
						self.parent.assertEqual(result.logits.shape   ,    (self.batch_size, self.num_choices)  )







			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.prepare_config_and_inputs()
						(
						    UpperCamelCase__				
						)       =							config_and_inputs
						UpperCamelCase__       =							{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """lengths""": input_lengths}
						return config, inputs_dict







@require_torch
class 							lowerCamelCase__						(				_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE       ,						_SCREAMING_SNAKE_CASE       ,						unittest.TestCase   ):






			"""simple docstring"""






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




			def 					snake_case__					(				self   ,    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 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case=False  ):

						'''simple docstring'''


						UpperCamelCase__       =							super()._prepare_for_class(snake_case   ,    snake_case   ,    return_labels=snake_case  )

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

						return inputs_dict




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMModelTester(self  )
						UpperCamelCase__       =							ConfigTester(self   ,    config_class=snake_case   ,    emb_dim=37  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						self.config_tester.run_common_tests()




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_model(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_lm_head(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_simple_qa(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_qa(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_sequence_classif(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_token_classif(*snake_case  )




			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_xlm_for_multiple_choice(*snake_case  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case=False   ,    snake_case=1  ):

						'''simple docstring'''


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

						for idx, iter_attentions in enumerate(snake_case  ):
									# adds PAD dummy token
									UpperCamelCase__       =							min_length + idx + 1
									UpperCamelCase__       =							min_length + idx + 1

									UpperCamelCase__       =							(
									    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(snake_case  )  )




			def 					snake_case__					(				self   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case   ,    snake_case=False   ,    snake_case=1  ):

						'''simple docstring'''


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

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







			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									UpperCamelCase__       =							XLMModel.from_pretrained(snake_case  )
									self.assertIsNotNone(snake_case  )







@require_torch
class 							lowerCamelCase__						(				unittest.TestCase   ):






			"""simple docstring"""






			@slow
			def 					snake_case__					(				self  ):

						'''simple docstring'''


						UpperCamelCase__       =							XLMWithLMHeadModel.from_pretrained("xlm-mlm-en-2048"  )
						model.to(snake_case  )
						UpperCamelCase__       =							torch.tensor([[14, 447]]   ,    dtype=torch.long   ,    device=snake_case  )  # the president
						UpperCamelCase__       =							[
						    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
						UpperCamelCase__       =							model.generate(snake_case   ,    do_sample=snake_case  )
						self.assertListEqual(output_ids[0].cpu().numpy().tolist()   ,    snake_case  )






'''simple docstring'''






from typing import Union

import fire
import torch
from tqdm import tqdm







def 							UpperCamelCase_     (       A__   :  str				,       A__   :  str = "cpu"				,       A__   :  Union[str, None] = None   ):
  '''simple docstring'''





  lowerCAmelCase_     : Dict     =  torch.load(A__				,       map_location=A__   )
  for k, v in tqdm(state_dict.items()   ):
    if not isinstance(A__				,       torch.Tensor   ):
      raise TypeError("""FP16 conversion only works on paths that are saved state dicts, like pytorch_model.bin"""   )
    lowerCAmelCase_     : List[str]     =  v.half()
  if save_path is None:  # overwrite src_path
    lowerCAmelCase_     : List[str]     =  src_path
  torch.save(A__				,       A__   )


if __name__ == "__main__":
  fire.Fire(convert)



