




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






		def __init__(    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=13						,		_SCREAMING_SNAKE_CASE=7						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=False						,		_SCREAMING_SNAKE_CASE=False						,		_SCREAMING_SNAKE_CASE=False						,		_SCREAMING_SNAKE_CASE=2						,		_SCREAMING_SNAKE_CASE=99						,		_SCREAMING_SNAKE_CASE=0						,		_SCREAMING_SNAKE_CASE=32						,		_SCREAMING_SNAKE_CASE=5						,		_SCREAMING_SNAKE_CASE=4						,		_SCREAMING_SNAKE_CASE=0.1						,		_SCREAMING_SNAKE_CASE=0.1						,		_SCREAMING_SNAKE_CASE=512						,		_SCREAMING_SNAKE_CASE=2						,		_SCREAMING_SNAKE_CASE=0.02						,		_SCREAMING_SNAKE_CASE=2						,		_SCREAMING_SNAKE_CASE=4						,		_SCREAMING_SNAKE_CASE="last"						,		_SCREAMING_SNAKE_CASE=True						,		_SCREAMING_SNAKE_CASE=None						,		_SCREAMING_SNAKE_CASE=0						,		):
								_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 				UpperCAmelCase (    self			):
								_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 				UpperCAmelCase (    self			):
								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 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			XLMModel(config=_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()
								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE						,		lengths=_SCREAMING_SNAKE_CASE						,		langs=_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE						,		langs=_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE			)
								self.parent.assertEqual(result.last_hidden_state.shape						,		(self.batch_size, self.seq_length, self.hidden_size)			)






		def 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			XLMWithLMHeadModel(_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()

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






		def 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			XLMForQuestionAnsweringSimple(_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()

								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE						,		start_positions=_SCREAMING_SNAKE_CASE						,		end_positions=_SCREAMING_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 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			XLMForQuestionAnswering(_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()

								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE			)

								_UpperCAmelCase      =			model(
								    _SCREAMING_SNAKE_CASE						,		start_positions=_SCREAMING_SNAKE_CASE						,		end_positions=_SCREAMING_SNAKE_CASE						,		cls_index=_SCREAMING_SNAKE_CASE						,		is_impossible=_SCREAMING_SNAKE_CASE						,		p_mask=_SCREAMING_SNAKE_CASE						,		)

								_UpperCAmelCase      =			model(
								    _SCREAMING_SNAKE_CASE						,		start_positions=_SCREAMING_SNAKE_CASE						,		end_positions=_SCREAMING_SNAKE_CASE						,		cls_index=_SCREAMING_SNAKE_CASE						,		is_impossible=_SCREAMING_SNAKE_CASE						,		)

								(_UpperCAmelCase						)      =			result_with_labels.to_tuple()

								_UpperCAmelCase      =			model(_SCREAMING_SNAKE_CASE						,		start_positions=_SCREAMING_SNAKE_CASE						,		end_positions=_SCREAMING_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 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			XLMForSequenceClassification(_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()

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






		def 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			self.num_labels
								_UpperCAmelCase      =			XLMForTokenClassification(_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_SNAKE_CASE			)
								model.eval()

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






		def 				UpperCAmelCase (    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						,		):
								_UpperCAmelCase      =			self.num_choices
								_UpperCAmelCase      =			XLMForMultipleChoice(config=_SCREAMING_SNAKE_CASE			)
								model.to(_SCREAMING_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(
								    _SCREAMING_SNAKE_CASE						,		attention_mask=_SCREAMING_SNAKE_CASE						,		token_type_ids=_SCREAMING_SNAKE_CASE						,		labels=_SCREAMING_SNAKE_CASE						,		)
								self.parent.assertEqual(result.logits.shape						,		(self.batch_size, self.num_choices)			)






		def 				UpperCAmelCase (    self			):
								_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       _A				(		_A       ,     _A       ,     _A       ,     unittest.TestCase     ):
		__a					    =       (
		    (
		        XLMModel,
		        XLMWithLMHeadModel,
		        XLMForQuestionAnswering,
		        XLMForSequenceClassification,
		        XLMForQuestionAnsweringSimple,
		        XLMForTokenClassification,
		        XLMForMultipleChoice,
		    )
		    if is_torch_available()
		    else ()
		)
		__a					    =       (
		    (XLMWithLMHeadModel,) if is_torch_available() else ()
		)  # TODO (PVP): Check other models whether language generation is also applicable
		__a					    =       (
		    {
		        """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 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			):
								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 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=False			):
								_UpperCAmelCase      =			super()._prepare_for_class(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		return_labels=_SCREAMING_SNAKE_CASE			)

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

								return inputs_dict






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			XLMModelTester(self			)
								_UpperCAmelCase      =			ConfigTester(self						,		config_class=_SCREAMING_SNAKE_CASE						,		emb_dim=37			)






		def 				UpperCAmelCase (    self			):
								self.config_tester.run_common_tests()






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_model(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_lm_head(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_simple_qa(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_qa(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_sequence_classif(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_token_classif(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_xlm_for_multiple_choice(*_SCREAMING_SNAKE_CASE			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=False						,		_SCREAMING_SNAKE_CASE=1			):
								self.assertIsInstance(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								self.assertListEqual(
								    [isinstance(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			) for iter_attentions in attentions]						,		[True] * len(_SCREAMING_SNAKE_CASE			)			)
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		(max_length - min_length) * num_beam_groups			)

								for idx, iter_attentions in enumerate(_SCREAMING_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(_SCREAMING_SNAKE_CASE			)			)






		def 				UpperCAmelCase (    self						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE=False						,		_SCREAMING_SNAKE_CASE=1			):
								self.assertIsInstance(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			)
								self.assertListEqual(
								    [isinstance(_SCREAMING_SNAKE_CASE						,		_SCREAMING_SNAKE_CASE			) for iter_hidden_states in hidden_states]						,		[True] * len(_SCREAMING_SNAKE_CASE			)						,		)
								self.assertEqual(len(_SCREAMING_SNAKE_CASE			)						,		(max_length - min_length) * num_beam_groups			)

								for idx, iter_hidden_states in enumerate(_SCREAMING_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(_SCREAMING_SNAKE_CASE			)						,		)
								pass






		@slow
		def 				UpperCAmelCase (    self			):
								for model_name in XLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_UpperCAmelCase      =			XLMModel.from_pretrained(_SCREAMING_SNAKE_CASE			)
														self.assertIsNotNone(_SCREAMING_SNAKE_CASE			)




@require_torch
class       _A				(		unittest.TestCase     ):






		@slow
		def 				UpperCAmelCase (    self			):
								_UpperCAmelCase      =			XLMWithLMHeadModel.from_pretrained("""xlm-mlm-en-2048"""			)
								model.to(_SCREAMING_SNAKE_CASE			)
								_UpperCAmelCase      =			torch.tensor([[14, 447]]						,		dtype=torch.long						,		device=_SCREAMING_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(_SCREAMING_SNAKE_CASE						,		do_sample=_SCREAMING_SNAKE_CASE			)
								self.assertListEqual(output_ids[0].cpu().numpy().tolist()						,		_SCREAMING_SNAKE_CASE			)





import jax.numpy as jnp

from ...utils import logging
from ..ta.modeling_flax_ta import FlaxTaEncoderModel, FlaxTaForConditionalGeneration, FlaxTaModel
from .configuration_mta import MTaConfig


a													=			logging.get_logger(__name__)

a													=			"T5Config"




def        _SCREAMING_SNAKE_CASE      (		snake_case			,     snake_case			,     snake_case       )						->				jnp.ndarray:
						_UpperCAmelCase      =			jnp.zeros_like(snake_case       )
						_UpperCAmelCase      =			shifted_input_ids.at[:, 1:].set(input_ids[:, :-1]       )
						_UpperCAmelCase      =			shifted_input_ids.at[:, 0].set(snake_case       )

						_UpperCAmelCase      =			jnp.where(shifted_input_ids == -1_0_0			,     snake_case			,     snake_case       )
						return shifted_input_ids




class       _A				(		__lowercase     ):
		__a					    =       """mt5"""
		__a					    =       MTaConfig




class       _A				(		__lowercase     ):
		__a					    =       """mt5"""
		__a					    =       MTaConfig




class       _A				(		__lowercase     ):
		__a					    =       """mt5"""
		__a					    =       MTaConfig