

"""simple docstring"""


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   :    Dict			,			_lowercase   :    Dict			,			_lowercase   :    Optional[Any]=13			,			_lowercase   :    List[Any]=7			,			_lowercase   :    Tuple=True			,			_lowercase   :    str=True			,			_lowercase   :    Dict=True			,			_lowercase   :    Optional[Any]=True			,			_lowercase   :    int=True			,			_lowercase   :    Dict=False			,			_lowercase   :    Optional[Any]=False			,			_lowercase   :    Optional[Any]=False			,			_lowercase   :    List[str]=2			,			_lowercase   :    List[Any]=99			,			_lowercase   :    Tuple=0			,			_lowercase   :    List[Any]=32			,			_lowercase   :    Any=5			,			_lowercase   :    Optional[int]=4			,			_lowercase   :    Tuple=0.1			,			_lowercase   :    Optional[Any]=0.1			,			_lowercase   :    Dict=512			,			_lowercase   :    List[str]=2			,			_lowercase   :    Dict=0.0_2			,			_lowercase   :    int=2			,			_lowercase   :    int=4			,			_lowercase   :    Optional[Any]="last"			,			_lowercase   :    List[Any]=True			,			_lowercase   :    Any=None			,			_lowercase   :    Union[str, Any]=0			,			):
								A					       = parent
								A					       = batch_size
								A					       = seq_length
								A					       = is_training
								A					       = use_input_lengths
								A					       = use_token_type_ids
								A					       = use_labels
								A					       = gelu_activation
								A					       = sinusoidal_embeddings
								A					       = causal
								A					       = asm
								A					       = n_langs
								A					       = vocab_size
								A					       = n_special
								A					       = hidden_size
								A					       = num_hidden_layers
								A					       = num_attention_heads
								A					       = hidden_dropout_prob
								A					       = attention_probs_dropout_prob
								A					       = max_position_embeddings
								A					       = type_sequence_label_size
								A					       = initializer_range
								A					       = num_labels
								A					       = num_choices
								A					       = summary_type
								A					       = use_proj
								A					       = scope
								A					       = bos_token_id







						def __a						(							self   :    Optional[int]		):
								A					       = ids_tensor([self.batch_size, self.seq_length]			,			self.vocab_size		)
								A					       = random_attention_mask([self.batch_size, self.seq_length]		)

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

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

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

								A					       = self.get_config()

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







						def __a						(							self   :    Any		):
								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 __a						(							self   :    Union[str, Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Tuple			,			_lowercase   :    Tuple			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Dict			,			_lowercase   :    Tuple			,			):
								A					       = XLMModel(config=_lowercase		)
								model.to(_lowercase		)
								model.eval()
								A					       = model(_lowercase			,			lengths=_lowercase			,			langs=_lowercase		)
								A					       = model(_lowercase			,			langs=_lowercase		)
								A					       = model(_lowercase		)
								self.parent.assertEqual(result.last_hidden_state.shape			,			(self.batch_size, self.seq_length, self.hidden_size)		)







						def __a						(							self   :    Optional[Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    Optional[int]			,			_lowercase   :    List[str]			,			_lowercase   :    int			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Dict			,			_lowercase   :    List[str]			,			):
								A					       = XLMWithLMHeadModel(_lowercase		)
								model.to(_lowercase		)
								model.eval()

								A					       = 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 __a						(							self   :    Any			,			_lowercase   :    List[Any]			,			_lowercase   :    Any			,			_lowercase   :    Tuple			,			_lowercase   :    int			,			_lowercase   :    int			,			_lowercase   :    int			,			_lowercase   :    Tuple			,			_lowercase   :    Dict			,			_lowercase   :    str			,			):
								A					       = XLMForQuestionAnsweringSimple(_lowercase		)
								model.to(_lowercase		)
								model.eval()

								A					       = model(_lowercase		)

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

								A					       = model(_lowercase		)

								A					       = model(
								    _lowercase			,			start_positions=_lowercase			,			end_positions=_lowercase			,			cls_index=_lowercase			,			is_impossible=_lowercase			,			p_mask=_lowercase			,			)

								A					       = model(
								    _lowercase			,			start_positions=_lowercase			,			end_positions=_lowercase			,			cls_index=_lowercase			,			is_impossible=_lowercase			,			)

								((A					)       , 					)					       = result_with_labels.to_tuple()

								A					       = model(_lowercase			,			start_positions=_lowercase			,			end_positions=_lowercase		)

								((A					)       , 					)					       = 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 __a						(							self   :    Any			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Any			,			_lowercase   :    Optional[Any]			,			_lowercase   :    List[Any]			,			_lowercase   :    int			,			_lowercase   :    Tuple			,			_lowercase   :    int			,			_lowercase   :    Optional[int]			,			_lowercase   :    Any			,			):
								A					       = XLMForSequenceClassification(_lowercase		)
								model.to(_lowercase		)
								model.eval()

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

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







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







						def __a						(							self   :    Dict		):
								A					       = self.prepare_config_and_inputs()
								(
								    (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , (
								    A					
								)       , 					
								)					       = config_and_inputs
								A					       = {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'lengths': input_lengths}
								return config, inputs_dict



@require_torch
class 			lowerCamelCase__						(   lowercase_							,    lowercase_							,    lowercase_							,    unittest.TestCase			):
						lowerCAmelCase       =     (
						    (
						        XLMModel,
						        XLMWithLMHeadModel,
						        XLMForQuestionAnswering,
						        XLMForSequenceClassification,
						        XLMForQuestionAnsweringSimple,
						        XLMForTokenClassification,
						        XLMForMultipleChoice,
						    )
						    if is_torch_available()
						    else ()
						)
						lowerCAmelCase       =     (
						    (XLMWithLMHeadModel,) if is_torch_available() else ()
						)  # TODO (PVP): Check other models whether language generation is also applicable
						lowerCAmelCase       =     (
						    {
						        '''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 __a						(							self   :    List[str]			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    Any			,			_lowercase   :    List[str]		):
								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   :    int			,			_lowercase   :    Tuple			,			_lowercase   :    Any			,			_lowercase   :    Union[str, Any]=False		):
								A					       = super()._prepare_for_class(_lowercase			,			_lowercase			,			return_labels=_lowercase		)

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

								return inputs_dict







						def __a						(							self   :    str		):
								A					       = XLMModelTester(self		)
								A					       = ConfigTester(self			,			config_class=_lowercase			,			emb_dim=37		)







						def __a						(							self   :    Any		):
								self.config_tester.run_common_tests()







						def __a						(							self   :    List[str]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_model(*_lowercase		)







						def __a						(							self   :    int		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_lm_head(*_lowercase		)







						def __a						(							self   :    str		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_simple_qa(*_lowercase		)







						def __a						(							self   :    str		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_qa(*_lowercase		)







						def __a						(							self   :    Union[str, Any]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_sequence_classif(*_lowercase		)







						def __a						(							self   :    Optional[int]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_token_classif(*_lowercase		)







						def __a						(							self   :    Union[str, Any]		):
								A					       = self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_for_multiple_choice(*_lowercase		)







						def __a						(							self   :    str			,			_lowercase   :    List[Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    Any			,			_lowercase   :    Optional[int]			,			_lowercase   :    List[str]			,			_lowercase   :    Dict=False			,			_lowercase   :    str=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
										A					       = min_length + idx + 1
										A					       = min_length + idx + 1

										A					       = (
										    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 __a						(							self   :    Optional[Any]			,			_lowercase   :    Optional[Any]			,			_lowercase   :    Any			,			_lowercase   :    Union[str, Any]			,			_lowercase   :    Optional[int]			,			_lowercase   :    str			,			_lowercase   :    int=False			,			_lowercase   :    Union[str, Any]=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
										A					       = min_length + idx + 1
										A					       = (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 __a						(							self   :    str		):
								for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										A					       = XLMModel.from_pretrained(_lowercase		)
										self.assertIsNotNone(_lowercase		)



@require_torch
class 			lowerCamelCase__						(   unittest.TestCase			):







						@slow
						def __a						(							self   :    Dict		):
								A					       = XLMWithLMHeadModel.from_pretrained('xlm-mlm-en-2048'		)
								model.to(_lowercase		)
								A					       = torch.tensor([[14, 447]]			,			dtype=torch.long			,			device=_lowercase		)  # the president
								A					       = [
								    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
								A					       = model.generate(_lowercase			,			do_sample=_lowercase		)
								self.assertListEqual(output_ids[0].cpu().numpy().tolist()			,			_lowercase		)









"""simple docstring"""


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCamelCase						:				Dict             =      {
    "configuration_megatron_bert": ["MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP", "MegatronBertConfig"],
}

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCamelCase						:				Any             =      [
	    "MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST",
	    "MegatronBertForCausalLM",
	    "MegatronBertForMaskedLM",
	    "MegatronBertForMultipleChoice",
	    "MegatronBertForNextSentencePrediction",
	    "MegatronBertForPreTraining",
	    "MegatronBertForQuestionAnswering",
	    "MegatronBertForSequenceClassification",
	    "MegatronBertForTokenClassification",
	    "MegatronBertModel",
	    "MegatronBertPreTrainedModel",
	]

if TYPE_CHECKING:
	from .configuration_megatron_bert import MEGATRON_BERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MegatronBertConfig

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_megatron_bert import (
		    MEGATRON_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MegatronBertForCausalLM,
		    MegatronBertForMaskedLM,
		    MegatronBertForMultipleChoice,
		    MegatronBertForNextSentencePrediction,
		    MegatronBertForPreTraining,
		    MegatronBertForQuestionAnswering,
		    MegatronBertForSequenceClassification,
		    MegatronBertForTokenClassification,
		    MegatronBertModel,
		    MegatronBertPreTrainedModel,
		)

else:
	import sys

	UpperCamelCase						:				Dict             =      _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)





