




import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
		import torch

		from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel




class 					snake_case__							:
				"""simple docstring"""



				def __init__(   self       :  List[Any]      ,   __lowerCamelCase       :  Optional[int]      ,   __lowerCamelCase       :  Optional[int]=99      ,   __lowerCamelCase       :  Any=13      ,   __lowerCamelCase       :  List[str]=7      ,   __lowerCamelCase       :  Union[str, Any]=9      ,   __lowerCamelCase       :  Union[str, Any]=True      ,   __lowerCamelCase       :  List[Any]=True      ,   __lowerCamelCase       :  Dict=False      ,   __lowerCamelCase       :  List[Any]=32      ,   __lowerCamelCase       :  List[Any]=5      ,   __lowerCamelCase       :  List[Any]=4      ,   __lowerCamelCase       :  Any=37      ,   __lowerCamelCase       :  Any=8      ,   __lowerCamelCase       :  Union[str, Any]=0.1      ,   __lowerCamelCase       :  Optional[Any]=0.002      ,   __lowerCamelCase       :  List[str]=1      ,   __lowerCamelCase       :  Any=0      ,   __lowerCamelCase       :  Union[str, Any]=0      ,   __lowerCamelCase       :  Union[str, Any]=None      ,   __lowerCamelCase       :  List[str]=None      ,   )      ->  Union[str, Any]:
							a					=    parent
							a					=    batch_size
							a					=    encoder_seq_length
							a					=    decoder_seq_length
							# For common tests
							a					=    self.decoder_seq_length
							a					=    is_training
							a					=    use_attention_mask
							a					=    use_labels
							a					=    vocab_size
							a					=    hidden_size
							a					=    num_hidden_layers
							a					=    num_attention_heads
							a					=    d_ff
							a					=    relative_attention_num_buckets
							a					=    dropout_rate
							a					=    initializer_factor
							a					=    eos_token_id
							a					=    pad_token_id
							a					=    decoder_start_token_id
							a					=    None
							a					=    decoder_layers





				def 							__UpperCAmelCase							(   self       :  List[Any]							)      ->  Union[str, Any]:
							return TaConfig.from_pretrained("google/umt5-base"							)





				def 							__UpperCAmelCase							(   self       :  List[Any]      ,   __lowerCamelCase       :  Union[str, Any]      ,   __lowerCamelCase       :  Dict      ,   __lowerCamelCase       :  Any      ,   __lowerCamelCase       :  List[str]=None      ,   __lowerCamelCase       :  Dict=None      ,   __lowerCamelCase       :  List[str]=None      ,   __lowerCamelCase       :  Union[str, Any]=None      ,   __lowerCamelCase       :  List[Any]=None      ,   )      ->  Optional[Any]:
							if attention_mask is None:
										a					=    input_ids.ne(config.pad_token_id							)
							if decoder_attention_mask is None:
										a					=    decoder_input_ids.ne(config.pad_token_id							)
							if head_mask is None:
										a					=    torch.ones(config.num_hidden_layers      ,   config.num_attention_heads      ,   device=__lowerCamelCase							)
							if decoder_head_mask is None:
										a					=    torch.ones(config.num_decoder_layers      ,   config.num_attention_heads      ,   device=__lowerCamelCase							)
							if cross_attn_head_mask is None:
										a					=    torch.ones(
										    config.num_decoder_layers      ,   config.num_attention_heads      ,   device=__lowerCamelCase							)
							return {
							    "input_ids": input_ids,
							    "decoder_input_ids": decoder_input_ids,
							    "attention_mask": attention_mask,
							    "decoder_attention_mask": decoder_attention_mask,
							    "head_mask": head_mask,
							    "decoder_head_mask": decoder_head_mask,
							    "cross_attn_head_mask": cross_attn_head_mask,
							}





				def 							__UpperCAmelCase							(   self       :  Optional[int]							)      ->  List[str]:
							a					=    ids_tensor([self.batch_size, self.encoder_seq_length]      ,   self.vocab_size							)
							a					=    ids_tensor([self.batch_size, self.decoder_seq_length]      ,   self.vocab_size							)

							# we need to clamp the input ids here to avoid having pad token in between
							# this is because for NllbMoe the position_ids are prepared such that
							# all pad tokens have pos id = 2 and rest are between 2..seq_length
							# and the seq_length here is seq_length - num_pad_tokens
							# but when using past, there is no way of knowing if the past input ids had
							# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
							# position_ids being off by num_pad_tokens in past input
							a					=    input_ids.clamp(self.pad_token_id + 1							)
							a					=    decoder_input_ids.clamp(self.pad_token_id + 1							)

							a					=    self.get_config()
							a					=    config.num_attention_heads
							a					=    self.prepare_inputs_dict(__lowerCamelCase      ,   __lowerCamelCase      ,   __lowerCamelCase							)
							return config, input_dict





				def 							__UpperCAmelCase							(   self       :  List[Any]							)      ->  Optional[Any]:
							a					,	a					=    self.prepare_config_and_inputs()
							return config, inputs_dict





				def 							__UpperCAmelCase							(   self       :  str							)      ->  int:
							return TaConfig(
							    vocab_size=1_66      ,   d_model=self.hidden_size      ,   d_ff=self.d_ff      ,   d_kv=self.hidden_size // self.num_attention_heads      ,   num_layers=self.num_hidden_layers      ,   num_decoder_layers=self.decoder_layers      ,   num_heads=self.num_attention_heads      ,   relative_attention_num_buckets=self.relative_attention_num_buckets      ,   dropout_rate=self.dropout_rate      ,   initializer_factor=self.initializer_factor      ,   eos_token_id=self.eos_token_id      ,   bos_token_id=self.pad_token_id      ,   pad_token_id=self.pad_token_id      ,   decoder_start_token_id=self.decoder_start_token_id      ,   )





				def 							__UpperCAmelCase							(   self       :  Union[str, Any]							)      ->  Dict:
							return TaConfig(
							    vocab_size=self.vocab_size      ,   d_model=self.hidden_size      ,   d_ff=self.d_ff      ,   d_kv=self.hidden_size // self.num_attention_heads      ,   num_layers=self.num_hidden_layers      ,   num_decoder_layers=self.decoder_layers      ,   num_heads=self.num_attention_heads      ,   relative_attention_num_buckets=self.relative_attention_num_buckets      ,   dropout_rate=self.dropout_rate      ,   initializer_factor=self.initializer_factor      ,   eos_token_id=self.eos_token_id      ,   bos_token_id=self.pad_token_id      ,   pad_token_id=self.pad_token_id      ,   decoder_start_token_id=self.decoder_start_token_id      ,   )





				def 							__UpperCAmelCase							(   self       :  Union[str, Any]      ,   __lowerCamelCase       :  str      ,   __lowerCamelCase       :  List[Any]      ,   __lowerCamelCase       :  Union[str, Any]      ,   __lowerCamelCase       :  Any      ,   __lowerCamelCase       :  List[Any]      ,   __lowerCamelCase       :  Dict      ,   )      ->  str:
							a					=    UMTaModel(config=__lowerCamelCase							)
							model.to(__lowerCamelCase							)
							model.eval()
							a					=    model(
							    input_ids=__lowerCamelCase      ,   decoder_input_ids=__lowerCamelCase      ,   attention_mask=__lowerCamelCase      ,   decoder_attention_mask=__lowerCamelCase      ,   )
							a					=    model(input_ids=__lowerCamelCase      ,   decoder_input_ids=__lowerCamelCase							)
							a					=    result.last_hidden_state
							a					=    result.past_key_values
							a					=    result.encoder_last_hidden_state

							self.parent.assertEqual(encoder_output.size()      ,   (self.batch_size, self.encoder_seq_length, self.hidden_size)							)
							self.parent.assertEqual(decoder_output.size()      ,   (self.batch_size, self.decoder_seq_length, self.hidden_size)							)
							# There should be `num_layers` key value embeddings stored in decoder_past
							self.parent.assertEqual(len(__lowerCamelCase							)      ,   config.num_layers							)
							# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
							self.parent.assertEqual(len(decoder_past[0]							)      ,   4							)





				def 							__UpperCAmelCase							(   self       :  Any      ,   __lowerCamelCase       :  Optional[int]      ,   __lowerCamelCase       :  List[Any]      ,   __lowerCamelCase       :  str      ,   __lowerCamelCase       :  List[str]      ,   __lowerCamelCase       :  int      ,   __lowerCamelCase       :  Tuple      ,   )      ->  List[Any]:
							a					=    UMTaModel(config=__lowerCamelCase							).get_decoder().to(__lowerCamelCase							).eval()
							# first forward pass
							a					=    model(__lowerCamelCase      ,   use_cache=__lowerCamelCase							)
							a					=    model(__lowerCamelCase							)
							a					=    model(__lowerCamelCase      ,   use_cache=__lowerCamelCase							)

							self.parent.assertTrue(len(__lowerCamelCase							) == len(__lowerCamelCase							)							)
							self.parent.assertTrue(len(__lowerCamelCase							) == len(__lowerCamelCase							) + 1							)

							a					,	a					=    outputs.to_tuple()

							# create hypothetical next token and extent to next_input_ids
							a					=    ids_tensor((self.batch_size, 1)      ,   config.vocab_size							)

							# append to next input_ids and
							a					=    torch.cat([input_ids, next_tokens]      ,   dim=-1							)

							a					=    model(__lowerCamelCase							)["last_hidden_state"]
							a					=    model(__lowerCamelCase      ,   past_key_values=__lowerCamelCase							)["last_hidden_state"]

							# select random slice
							a					=    ids_tensor((1,)      ,   output_from_past.shape[-1]							).item()
							a					=    output_from_no_past[:, -1, random_slice_idx].detach()
							a					=    output_from_past[:, 0, random_slice_idx].detach()

							# test that outputs are equal for slice
							self.parent.assertTrue(torch.allclose(__lowerCamelCase      ,   __lowerCamelCase      ,   atol=1e-3							)							)





				def 							__UpperCAmelCase							(   self       :  int      ,   __lowerCamelCase       :  str      ,   __lowerCamelCase       :  int      ,   )      ->  List[str]:
							a					=    UMTaModel(config=__lowerCamelCase							).to(__lowerCamelCase							).half().eval()
							a					=    model(**__lowerCamelCase							)["last_hidden_state"]
							self.parent.assertFalse(torch.isnan(__lowerCamelCase							).any().item()							)




@require_torch
class 					snake_case__							(_UpperCamelCase					, _UpperCamelCase					, _UpperCamelCase					, unittest.TestCase     ):
				"""simple docstring"""



				SCREAMING_SNAKE_CASE_ :							Union[str, Any]											=      (
				    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
				)
				SCREAMING_SNAKE_CASE_ :							int											=      (UMTaForConditionalGeneration,) if is_torch_available() else ()
				SCREAMING_SNAKE_CASE_ :							str											=      (
				    {
				        """conversational""": UMTaForConditionalGeneration,
				        """feature-extraction""": UMTaModel,
				        """summarization""": UMTaForConditionalGeneration,
				        """text2text-generation""": UMTaForConditionalGeneration,
				        """translation""": UMTaForConditionalGeneration,
				        """question-answering""": UMTaForQuestionAnswering,
				    }
				    if is_torch_available()
				    else {}
				)
				SCREAMING_SNAKE_CASE_ :							Optional[int]											=      True
				SCREAMING_SNAKE_CASE_ :							Union[str, Any]											=      False
				SCREAMING_SNAKE_CASE_ :							List[Any]											=      False
				SCREAMING_SNAKE_CASE_ :							List[Any]											=      True
				SCREAMING_SNAKE_CASE_ :							Union[str, Any]											=      True
				# The small UMT5 model needs higher percentages for CPU/MP tests
				SCREAMING_SNAKE_CASE_ :							Tuple											=      [0.8, 0.9]





				def 							__UpperCAmelCase							(   self       :  Dict							)      ->  int:
							a					=    UMTaModelTester(self							)





				@unittest.skip("Test has a segmentation fault on torch 1.8.0"							)
				def 							__UpperCAmelCase							(   self       :  int							)      ->  int:
							a					=    self.model_tester.prepare_config_and_inputs()
							a					=    UMTaModel(config_and_inputs[0]							).to(__lowerCamelCase							)
							with tempfile.TemporaryDirectory() as tmpdirname:
										torch.onnx.export(
										    __lowerCamelCase      ,   (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])      ,   f"""{tmpdirname}/t5_test.onnx"""      ,   export_params=__lowerCamelCase      ,   opset_version=9      ,   input_names=["input_ids", "decoder_input_ids"]      ,   )





				@unittest.skipIf(torch_device == "cpu"      ,   "Cant do half precision"							)
				def 							__UpperCAmelCase							(   self       :  int							)      ->  Optional[int]:
							a					=    self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model_fpaa_forward(*__lowerCamelCase							)





				def 							__UpperCAmelCase							(   self       :  int							)      ->  Optional[int]:
							a					=    ["encoder_attentions", "decoder_attentions", "cross_attentions"]
							a					=    self.model_tester.prepare_config_and_inputs()
							a					=    config_and_inputs[0]
							a					=    UMTaForConditionalGeneration(__lowerCamelCase							).eval()
							model.to(__lowerCamelCase							)

							a					=    {
							    "head_mask": torch.zeros(config.num_layers      ,   config.num_heads      ,   device=__lowerCamelCase							),
							    "decoder_head_mask": torch.zeros(config.num_decoder_layers      ,   config.num_heads      ,   device=__lowerCamelCase							),
							    "cross_attn_head_mask": torch.zeros(config.num_decoder_layers      ,   config.num_heads      ,   device=__lowerCamelCase							),
							}

							for attn_name, (name, mask) in zip(__lowerCamelCase      ,   head_masking.items()							):
										a					=    {name: mask}
										# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
										if name == "head_mask":
													a					=    torch.ones(
													    config.num_decoder_layers      ,   config.num_heads      ,   device=__lowerCamelCase							)

										a					=    model.generate(
										    config_and_inputs[1]["input_ids"]      ,   num_beams=1      ,   max_length=3      ,   output_attentions=__lowerCamelCase      ,   return_dict_in_generate=__lowerCamelCase      ,   **__lowerCamelCase      ,   )
										# We check the state of decoder_attentions and cross_attentions just from the last step
										a					=    out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
										self.assertEqual(sum([w.sum().item() for w in attn_weights]							)      ,   0.0							)





				@unittest.skip("Does not work on the tiny model as we keep hitting edge cases."							)
				def 							__UpperCAmelCase							(   self       :  Optional[Any]							)      ->  int:
							pass




@require_torch
@require_sentencepiece
@require_tokenizers
class 					snake_case__							(unittest.TestCase     ):
				"""simple docstring"""



				@slow
				@unittest.skip(
				    "Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged"							)
				def 							__UpperCAmelCase							(   self       :  str							)      ->  str:
							a					=    UMTaForConditionalGeneration.from_pretrained("google/umt5-small"      ,   return_dict=__lowerCamelCase							).to(__lowerCamelCase							)
							a					=    AutoTokenizer.from_pretrained("google/umt5-small"      ,   use_fast=__lowerCamelCase      ,   legacy=__lowerCamelCase							)
							a					=    [
							    "Bonjour monsieur <extra_id_0> bien <extra_id_1>.",
							    "No se como puedo <extra_id_0>.",
							    "This is the reason why we <extra_id_0> them.",
							    "The <extra_id_0> walks in <extra_id_1>, seats",
							    "A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.",
							]
							a					=    tokenizer(__lowerCamelCase      ,   return_tensors="pt"      ,   padding=__lowerCamelCase							).input_ids
							# fmt: off
							a					=    torch.tensor(
							    [
							        [ 3_85_30, 21_07_03, 25_62_99, 14_10, 25_62_98, 2_74, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   8_26, 3_21, 6_71, 2_59_22, 25_62_99, 2_74, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [  14_60, 3_39, 3_12, 1_90_14, 1_06_20, 7_58, 25_62_99, 23_55,2_74, 1, 0, 0, 0, 0, 0, 0,0, 0],
							        [   5_17, 25_62_99, 1_48_69, 2_81, 3_01, 25_62_98, 2_75, 11_99_83,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
							        [   3_20, 25_62_99, 1_48_69, 2_81, 22_34, 2_89, 22_75, 3_33,6_13_91, 2_89, 25_62_98, 5_43, 25_62_97, 16_87_14, 3_29, 25_62_96,2_74, 1],
							    ]							)
							# fmt: on
							torch.testing.assert_allclose(__lowerCamelCase      ,   __lowerCamelCase							)

							a					=    model.generate(input_ids.to(__lowerCamelCase							)							)
							a					=    [
							    "<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>",
							    "<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							    "<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>",
							]
							a					=    tokenizer.batch_decode(__lowerCamelCase							)
							self.assertEqual(__lowerCamelCase      ,   __lowerCamelCase							)












from typing import List, Optional, TypeVar

from .arrow_dataset import Dataset, _concatenate_map_style_datasets, _interleave_map_style_datasets
from .dataset_dict import DatasetDict, IterableDatasetDict
from .info import DatasetInfo
from .iterable_dataset import IterableDataset, _concatenate_iterable_datasets, _interleave_iterable_datasets
from .splits import NamedSplit
from .utils import logging
from .utils.py_utils import Literal


__lowerCAmelCase	:       List[Any]     =	logging.get_logger(__name__)


__lowerCAmelCase	:       List[Any]     =	TypeVar('DatasetType', Dataset, IterableDataset)




def     __magic_name__       (				A	:      List[DatasetType],						A	:      Optional[List[float]] = None,						A	:      Optional[int] = None,						A	:      Optional[DatasetInfo] = None,						A	:      Optional[NamedSplit] = None,						A	:      Literal["first_exhausted", "all_exhausted"] = "first_exhausted",						):






			'''simple docstring'''







			from .arrow_dataset import Dataset
			from .iterable_dataset import IterableDataset

			if not datasets:
						raise ValueError("Unable to interleave an empty list of datasets."     )
			for i, dataset in enumerate(A     ):
						if not isinstance(A,						(Dataset, IterableDataset)     ):
									if isinstance(A,						(DatasetDict, IterableDatasetDict)     ):
												if not dataset:
															raise ValueError(
															    F"""Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} """
															    "is an empty dataset dictionary."     )
												raise ValueError(
												    F"""Dataset at position {i} has at least one split: {list(A     )}\n"""
												    F"""Please pick one to interleave with the other datasets, for example: dataset['{next(iter(A     )     )}']"""     )
									raise ValueError(
									    F"""Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(A     ).__name__}."""     )
						if i == 0:
									a					,	a					=    (
									    (Dataset, IterableDataset) if isinstance(A,						A     ) else (IterableDataset, Dataset)
									)
						elif not isinstance(A,						A     ):
									raise ValueError(
									    F"""Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."""     )
			if stopping_strategy not in ["first_exhausted", "all_exhausted"]:
						raise ValueError(F"""{stopping_strategy} is not supported. Please enter a valid stopping_strategy."""     )
			if dataset_type is Dataset:
						return _interleave_map_style_datasets(
						    A,						A,						A,						info=A,						split=A,						stopping_strategy=A     )
			else:
						return _interleave_iterable_datasets(
						    A,						A,						A,						info=A,						split=A,						stopping_strategy=A     )






def     __magic_name__       (				A	:      List[DatasetType],						A	:      Optional[DatasetInfo] = None,						A	:      Optional[NamedSplit] = None,						A	:      int = 0,						):






			'''simple docstring'''







			if not dsets:
						raise ValueError("Unable to concatenate an empty list of datasets."     )
			for i, dataset in enumerate(A     ):
						if not isinstance(A,						(Dataset, IterableDataset)     ):
									if isinstance(A,						(DatasetDict, IterableDatasetDict)     ):
												if not dataset:
															raise ValueError(
															    F"""Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} """
															    "is an empty dataset dictionary."     )
												raise ValueError(
												    F"""Dataset at position {i} has at least one split: {list(A     )}\n"""
												    F"""Please pick one to interleave with the other datasets, for example: dataset['{next(iter(A     )     )}']"""     )
									raise ValueError(
									    F"""Expected a list of Dataset objects or a list of IterableDataset objects, but element at position {i} is a {type(A     ).__name__}."""     )
						if i == 0:
									a					,	a					=    (
									    (Dataset, IterableDataset) if isinstance(A,						A     ) else (IterableDataset, Dataset)
									)
						elif not isinstance(A,						A     ):
									raise ValueError(
									    F"""Unable to interleave a {dataset_type.__name__} (at position 0) with a {other_type.__name__} (at position {i}). Expected a list of Dataset objects or a list of IterableDataset objects."""     )
			if dataset_type is Dataset:
						return _concatenate_map_style_datasets(A,						info=A,						split=A,						axis=A     )
			else:
						return _concatenate_iterable_datasets(A,						info=A,						split=A,						axis=A     )






