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





			def __init__(       self	:	List[Any]  ,  lowercase__	:	Any  ,  lowercase__	:	str=13  ,  lowercase__	:	List[str]=7  ,  lowercase__	:	List[str]=True  ,  lowercase__	:	Any=True  ,  lowercase__	:	Tuple=True  ,  lowercase__	:	Union[str, Any]=True  ,  lowercase__	:	Tuple=True  ,  lowercase__	:	Dict=False  ,  lowercase__	:	List[Any]=False  ,  lowercase__	:	Any=False  ,  lowercase__	:	Dict=2  ,  lowercase__	:	Any=99  ,  lowercase__	:	Tuple=0  ,  lowercase__	:	Optional[Any]=32  ,  lowercase__	:	Union[str, Any]=5  ,  lowercase__	:	Dict=4  ,  lowercase__	:	str=0.1  ,  lowercase__	:	Tuple=0.1  ,  lowercase__	:	List[str]=5_12  ,  lowercase__	:	Dict=2  ,  lowercase__	:	Dict=0.0_2  ,  lowercase__	:	List[Any]=2  ,  lowercase__	:	Tuple=4  ,  lowercase__	:	Optional[Any]="last"  ,  lowercase__	:	Optional[Any]=True  ,  lowercase__	:	List[str]=None  ,  lowercase__	:	Union[str, Any]=0  ,  ):
										_lowerCAmelCase							=							parent
										_lowerCAmelCase							=							batch_size
										_lowerCAmelCase							=							seq_length
										_lowerCAmelCase							=							is_training
										_lowerCAmelCase							=							use_input_lengths
										_lowerCAmelCase							=							use_token_type_ids
										_lowerCAmelCase							=							use_labels
										_lowerCAmelCase							=							gelu_activation
										_lowerCAmelCase							=							sinusoidal_embeddings
										_lowerCAmelCase							=							causal
										_lowerCAmelCase							=							asm
										_lowerCAmelCase							=							n_langs
										_lowerCAmelCase							=							vocab_size
										_lowerCAmelCase							=							n_special
										_lowerCAmelCase							=							hidden_size
										_lowerCAmelCase							=							num_hidden_layers
										_lowerCAmelCase							=							num_attention_heads
										_lowerCAmelCase							=							hidden_dropout_prob
										_lowerCAmelCase							=							attention_probs_dropout_prob
										_lowerCAmelCase							=							max_position_embeddings
										_lowerCAmelCase							=							type_sequence_label_size
										_lowerCAmelCase							=							initializer_range
										_lowerCAmelCase							=							num_labels
										_lowerCAmelCase							=							num_choices
										_lowerCAmelCase							=							summary_type
										_lowerCAmelCase							=							use_proj
										_lowerCAmelCase							=							scope
										_lowerCAmelCase							=							bos_token_id





			def   SCREAMING_SNAKE_CASE__     (       self	:	str						):
										_lowerCAmelCase							=							ids_tensor([self.batch_size, self.seq_length]  ,  self.vocab_size						)
										_lowerCAmelCase							=							random_attention_mask([self.batch_size, self.seq_length]						)

										_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						).float()
																	_lowerCAmelCase							=							ids_tensor([self.batch_size]  ,  self.num_choices						)

										_lowerCAmelCase							=							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	:	int						):
										return XLMConfig(
										    vocab_size=self.vocab_size  ,  n_special=self.n_special  ,  emb_dim=self.hidden_size  ,  n_layers=self.num_hidden_layers  ,  n_heads=self.num_attention_heads  ,  dropout=self.hidden_dropout_prob  ,  attention_dropout=self.attention_probs_dropout_prob  ,  gelu_activation=self.gelu_activation  ,  sinusoidal_embeddings=self.sinusoidal_embeddings  ,  asm=self.asm  ,  causal=self.causal  ,  n_langs=self.n_langs  ,  max_position_embeddings=self.max_position_embeddings  ,  initializer_range=self.initializer_range  ,  summary_type=self.summary_type  ,  use_proj=self.use_proj  ,  num_labels=self.num_labels  ,  bos_token_id=self.bos_token_id  ,  )





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	Any  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	str  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Optional[Any]  ,  ):
										_lowerCAmelCase							=							XLMModel(config=lowercase__						)
										model.to(lowercase__						)
										model.eval()
										_lowerCAmelCase							=							model(lowercase__  ,  lengths=lowercase__  ,  langs=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  langs=lowercase__						)
										_lowerCAmelCase							=							model(lowercase__						)
										self.parent.assertEqual(result.last_hidden_state.shape  ,  (self.batch_size, self.seq_length, self.hidden_size)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[Any]  ,  lowercase__	:	Any  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Tuple  ,  lowercase__	:	int  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Any  ,  lowercase__	:	str  ,  lowercase__	:	int  ,  ):
										_lowerCAmelCase							=							XLMWithLMHeadModel(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__  ,  token_type_ids=lowercase__  ,  labels=lowercase__						)
										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	:	str  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Dict  ,  lowercase__	:	str  ,  lowercase__	:	Tuple  ,  lowercase__	:	List[str]  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	str  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]  ,  ):
										_lowerCAmelCase							=							XLMForQuestionAnsweringSimple(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)

										_lowerCAmelCase							=							model(lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__						)
										_lowerCAmelCase							=							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]  ,  lowercase__	:	List[str]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	Tuple  ,  lowercase__	:	List[str]  ,  lowercase__	:	Dict  ,  lowercase__	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  ):
										_lowerCAmelCase							=							XLMForQuestionAnswering(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)

										_lowerCAmelCase							=							model(
										    lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__  ,  cls_index=lowercase__  ,  is_impossible=lowercase__  ,  p_mask=lowercase__  ,  )

										_lowerCAmelCase							=							model(
										    lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__  ,  cls_index=lowercase__  ,  is_impossible=lowercase__  ,  )

										((_lowerCAmelCase)			,		)							=							result_with_labels.to_tuple()

										_lowerCAmelCase							=							model(lowercase__  ,  start_positions=lowercase__  ,  end_positions=lowercase__						)

										((_lowerCAmelCase)			,		)							=							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	:	Dict  ,  lowercase__	:	int  ,  lowercase__	:	List[str]  ,  lowercase__	:	Any  ,  lowercase__	:	Any  ,  lowercase__	:	int  ,  lowercase__	:	str  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	Optional[int]  ,  lowercase__	:	Optional[int]  ,  ):
										_lowerCAmelCase							=							XLMForSequenceClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()

										_lowerCAmelCase							=							model(lowercase__						)
										_lowerCAmelCase							=							model(lowercase__  ,  labels=lowercase__						)
										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	:	Optional[int]  ,  lowercase__	:	Dict  ,  lowercase__	:	int  ,  lowercase__	:	Dict  ,  lowercase__	:	int  ,  lowercase__	:	str  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	List[str]  ,  lowercase__	:	Optional[Any]  ,  lowercase__	:	Any  ,  ):
										_lowerCAmelCase							=							self.num_labels
										_lowerCAmelCase							=							XLMForTokenClassification(lowercase__						)
										model.to(lowercase__						)
										model.eval()

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





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





			def   SCREAMING_SNAKE_CASE__     (       self	:	Dict						):
										_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, 'lengths': input_lengths}
										return config, inputs_dict







@require_torch
class        lowerCamelCase__				(   UpperCAmelCase       ,UpperCAmelCase       ,UpperCAmelCase       ,unittest.TestCase				):
			UpperCamelCase__    							=(
			    (
			        XLMModel,
			        XLMWithLMHeadModel,
			        XLMForQuestionAnswering,
			        XLMForSequenceClassification,
			        XLMForQuestionAnsweringSimple,
			        XLMForTokenClassification,
			        XLMForMultipleChoice,
			    )
			    if is_torch_available()
			    else ()
			)
			UpperCamelCase__    							=(
			    (XLMWithLMHeadModel,) if is_torch_available() else ()
			)  # TODO (PVP): Check other models whether language generation is also applicable
			UpperCamelCase__    							=(
			    {
			        "feature-extraction": XLMModel,
			        "fill-mask": XLMWithLMHeadModel,
			        "question-answering": XLMForQuestionAnsweringSimple,
			        "text-classification": XLMForSequenceClassification,
			        "text-generation": XLMWithLMHeadModel,
			        "token-classification": XLMForTokenClassification,
			        "zero-shot": XLMForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int  ,  lowercase__	:	Any  ,  lowercase__	:	int  ,  lowercase__	:	List[str]  ,  lowercase__	:	str  ,  lowercase__	:	Tuple						):
										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  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	Tuple=False						):
										_lowerCAmelCase							=							super()._prepare_for_class(lowercase__  ,  lowercase__  ,  return_labels=lowercase__						)

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

										return inputs_dict





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							XLMModelTester(self						)
										_lowerCAmelCase							=							ConfigTester(self  ,  config_class=lowercase__  ,  emb_dim=37						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										self.config_tester.run_common_tests()





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_model(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	int						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_lm_head(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_simple_qa(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_qa(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_sequence_classif(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_token_classif(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Optional[int]						):
										_lowerCAmelCase							=							self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_xlm_for_multiple_choice(*lowercase__						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	List[str]  ,  lowercase__	:	Dict  ,  lowercase__	:	List[Any]  ,  lowercase__	:	List[Any]  ,  lowercase__	:	str  ,  lowercase__	:	int  ,  lowercase__	:	str=False  ,  lowercase__	:	Dict=1						):
										self.assertIsInstance(lowercase__  ,  lowercase__						)
										self.assertListEqual(
										    [isinstance(lowercase__  ,  lowercase__						) for iter_attentions in attentions]  ,  [True] * len(lowercase__						)						)
										self.assertEqual(len(lowercase__						)  ,  (max_length - min_length) * num_beam_groups						)

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

																	_lowerCAmelCase							=							(
																	    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(lowercase__						)						)





			def   SCREAMING_SNAKE_CASE__     (       self	:	Any  ,  lowercase__	:	Dict  ,  lowercase__	:	Union[str, Any]  ,  lowercase__	:	int  ,  lowercase__	:	Tuple  ,  lowercase__	:	List[str]  ,  lowercase__	:	Any=False  ,  lowercase__	:	str=1						):
										self.assertIsInstance(lowercase__  ,  lowercase__						)
										self.assertListEqual(
										    [isinstance(lowercase__  ,  lowercase__						) for iter_hidden_states in hidden_states]  ,  [True] * len(lowercase__						)  ,  )
										self.assertEqual(len(lowercase__						)  ,  (max_length - min_length) * num_beam_groups						)

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





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	Any						):
										for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	_lowerCAmelCase							=							XLMModel.from_pretrained(lowercase__						)
																	self.assertIsNotNone(lowercase__						)







@require_torch
class        lowerCamelCase__				(   unittest.TestCase				):





			@slow
			def   SCREAMING_SNAKE_CASE__     (       self	:	List[Any]						):
										_lowerCAmelCase							=							XLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'						)
										model.to(lowercase__						)
										_lowerCAmelCase							=							torch.tensor([[14, 4_47]]  ,  dtype=torch.long  ,  device=lowercase__						)  # the president
										_lowerCAmelCase							=							[
										    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							=							model.generate(lowercase__  ,  do_sample=lowercase__						)
										self.assertListEqual(output_ids[0].cpu().numpy().tolist()  ,  lowercase__						)







def _lowerCamelCase  (		snake_case  ):
							assert (
							    isinstance(snake_case     , snake_case  ) and number_of_steps > 0
							), F'number_of_steps needs to be positive integer, your input {number_of_steps}'
							if number_of_steps == 1:
														return 1
							_lowerCAmelCase			,		_lowerCAmelCase							=							1, 1
							for _ in range(number_of_steps - 1  ):
														_lowerCAmelCase			,		_lowerCAmelCase							=							current + previous, current
							return current


if __name__ == "__main__":
				import doctest

				doctest.testmod()






