





'''simple docstring'''

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






				'''simple docstring'''





				def __init__(				self       ,					_a       ,					_a=99       ,					_a=13       ,					_a=7       ,					_a=9       ,					_a=True       ,					_a=True       ,					_a=False       ,					_a=32       ,					_a=5       ,					_a=4       ,					_a=37       ,					_a=8       ,					_a=0.1       ,					_a=0.002       ,					_a=1       ,					_a=0       ,					_a=0       ,					_a=None       ,					_a=None       ,					):
											"""simple docstring"""
											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        lowercase__   (				self    ):
											"""simple docstring"""
											return TaConfig.from_pretrained('google/umt5-base'    )






				def        lowercase__   (				self       ,					_a       ,					_a       ,					_a       ,					_a=None       ,					_a=None       ,					_a=None       ,					_a=None       ,					_a=None       ,					):
											"""simple docstring"""
											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=__A    )
											if decoder_head_mask is None:
																		a__            =							torch.ones(config.num_decoder_layers       ,					config.num_attention_heads       ,					device=__A    )
											if cross_attn_head_mask is None:
																		a__            =							torch.ones(
																		    config.num_decoder_layers       ,					config.num_attention_heads       ,					device=__A    )
											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        lowercase__   (				self    ):
											"""simple docstring"""
											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(__A       ,					__A       ,					__A    )
											return config, input_dict






				def        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							self.prepare_config_and_inputs()
											return config, inputs_dict






				def        lowercase__   (				self    ):
											"""simple docstring"""
											return TaConfig(
											    vocab_size=166       ,					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        lowercase__   (				self    ):
											"""simple docstring"""
											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        lowercase__   (				self       ,					_a       ,					_a       ,					_a       ,					_a       ,					_a       ,					_a       ,					):
											"""simple docstring"""
											a__            =							UMTaModel(config=__A    )
											model.to(__A    )
											model.eval()
											a__            =							model(
											    input_ids=__A       ,					decoder_input_ids=__A       ,					attention_mask=__A       ,					decoder_attention_mask=__A       ,					)
											a__            =							model(input_ids=__A       ,					decoder_input_ids=__A    )
											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(__A    )       ,					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        lowercase__   (				self       ,					_a       ,					_a       ,					_a       ,					_a       ,					_a       ,					_a       ,					):
											"""simple docstring"""
											a__            =							UMTaModel(config=__A    ).get_decoder().to(__A    ).eval()
											# first forward pass
											a__            =							model(__A       ,					use_cache=__A    )
											a__            =							model(__A    )
											a__            =							model(__A       ,					use_cache=__A    )

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

											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(__A    )['''last_hidden_state''']
											a__            =							model(__A       ,					past_key_values=__A    )['''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(__A       ,					__A       ,					atol=1e-3    )    )






				def        lowercase__   (				self       ,					_a       ,					_a       ,					):
											"""simple docstring"""
											a__            =							UMTaModel(config=__A    ).to(__A    ).half().eval()
											a__            =							model(**__A    )['''last_hidden_state''']
											self.parent.assertFalse(torch.isnan(__A    ).any().item()    )






@require_torch
class       _UpperCamelCase      (			_A		,    _A		,    _A		,    unittest.TestCase			):






				'''simple docstring'''





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






				def        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							UMTaModelTester(self    )






				@unittest.skip('Test has a segmentation fault on torch 1.8.0'    )
				def        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							self.model_tester.prepare_config_and_inputs()
											a__            =							UMTaModel(config_and_inputs[0]    ).to(__A    )
											with tempfile.TemporaryDirectory() as tmpdirname:
																		torch.onnx.export(
																		    __A       ,					(config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])       ,					F'''{tmpdirname}/t5_test.onnx'''       ,					export_params=__A       ,					opset_version=9       ,					input_names=['input_ids', 'decoder_input_ids']       ,					)






				@unittest.skipIf(torch_device == 'cpu'       ,					'Cant do half precision'    )
				def        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model_fpaa_forward(*__A    )






				def        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
											a__            =							self.model_tester.prepare_config_and_inputs()
											a__            =							config_and_inputs[0]
											a__            =							UMTaForConditionalGeneration(__A    ).eval()
											model.to(__A    )

											a__            =							{
											    '''head_mask''': torch.zeros(config.num_layers       ,					config.num_heads       ,					device=__A    ),
											    '''decoder_head_mask''': torch.zeros(config.num_decoder_layers       ,					config.num_heads       ,					device=__A    ),
											    '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers       ,					config.num_heads       ,					device=__A    ),
											}

											for attn_name, (name, mask) in zip(__A       ,					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=__A    )

																		a__            =							model.generate(
																		    config_and_inputs[1]['input_ids']       ,					num_beams=1       ,					max_length=3       ,					output_attentions=__A       ,					return_dict_in_generate=__A       ,					**__A       ,					)
																		# 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        lowercase__   (				self    ):
											"""simple docstring"""
											pass






@require_torch
@require_sentencepiece
@require_tokenizers
class       _UpperCamelCase      (			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        lowercase__   (				self    ):
											"""simple docstring"""
											a__            =							UMTaForConditionalGeneration.from_pretrained('google/umt5-small'       ,					return_dict=__A    ).to(__A    )
											a__            =							AutoTokenizer.from_pretrained('google/umt5-small'       ,					use_fast=__A       ,					legacy=__A    )
											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(__A       ,					return_tensors='pt'       ,					padding=__A    ).input_ids
											# fmt: off
											a__            =							torch.tensor(
											    [
											        [ 3_8530, 21_0703, 25_6299, 1410, 25_6298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
											        [   826, 321, 671, 2_5922, 25_6299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
											        [  1460, 339, 312, 1_9014, 1_0620, 758, 25_6299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
											        [   517, 25_6299, 1_4869, 281, 301, 25_6298, 275, 11_9983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
											        [   320, 25_6299, 1_4869, 281, 2234, 289, 2275, 333,6_1391, 289, 25_6298, 543, 25_6297, 16_8714, 329, 25_6296,274, 1],
											    ]    )
											# fmt: on
											torch.testing.assert_allclose(__A       ,					__A    )

											a__            =							model.generate(input_ids.to(__A    )    )
											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(__A    )
											self.assertEqual(__A       ,					__A    )







from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


_lowercase											=    {
    """configuration_falcon""": ["""FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP""", """FalconConfig"""],
}

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 _lowercase											=    [
     """FALCON_PRETRAINED_MODEL_ARCHIVE_LIST""",
     """FalconForCausalLM""",
     """FalconModel""",
     """FalconPreTrainedModel""",
     """FalconForSequenceClassification""",
     """FalconForTokenClassification""",
     """FalconForQuestionAnswering""",
 ]


if TYPE_CHECKING:
 from .configuration_falcon import FALCON_PRETRAINED_CONFIG_ARCHIVE_MAP, FalconConfig

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_falcon import (
      FALCON_PRETRAINED_MODEL_ARCHIVE_LIST,
      FalconForCausalLM,
      FalconForQuestionAnswering,
      FalconForSequenceClassification,
      FalconForTokenClassification,
      FalconModel,
      FalconPreTrainedModel,
  )


else:
 import sys

 _lowercase											=    _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)
