


import doctest
import glob
import importlib
import inspect
import os
import re
from contextlib import contextmanager
from functools import wraps
from unittest.mock import patch

import numpy as np
import pytest
from absl.testing import parameterized

import datasets
from datasets import load_metric

from .utils import for_all_test_methods, local, slow


# mark all tests as integration
UpperCamelCase								=	pytest.mark.integration


UpperCamelCase								=	{"""comet"""}
UpperCamelCase								=	importlib.util.find_spec("""fairseq""") is not None

UpperCamelCase								=	{"""code_eval"""}
UpperCamelCase								=	os.name == """nt"""

UpperCamelCase								=	{"""bertscore""", """frugalscore""", """perplexity"""}
UpperCamelCase								=	importlib.util.find_spec("""transformers""") is not None


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  @wraps(SCREAMING_SNAKE_CASE       )
  def wrapper(self ,			SCREAMING_SNAKE_CASE       ):
    if not _has_fairseq and metric_name in REQUIRE_FAIRSEQ:
      self.skipTest('''"test requires Fairseq"'''       )
    else:
      test_case(self ,			SCREAMING_SNAKE_CASE       )

  return wrapper


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  @wraps(SCREAMING_SNAKE_CASE       )
  def wrapper(self ,			SCREAMING_SNAKE_CASE       ):
    if not _has_transformers and metric_name in REQUIRE_TRANSFORMERS:
      self.skipTest('''"test requires transformers"'''       )
    else:
      test_case(self ,			SCREAMING_SNAKE_CASE       )

  return wrapper


def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  @wraps(SCREAMING_SNAKE_CASE       )
  def wrapper(self ,			SCREAMING_SNAKE_CASE       ):
    if _on_windows and metric_name in UNSUPPORTED_ON_WINDOWS:
      self.skipTest('''"test not supported on Windows"'''       )
    else:
      test_case(self ,			SCREAMING_SNAKE_CASE       )

  return wrapper


def   _SCREAMING_SNAKE_CASE	( ):
  A_      :				List[str]		     =					[metric_dir.split(os.sep       )[-2] for metric_dir in glob.glob('''./metrics/*/'''       )]
  return [{"testcase_name": x, "metric_name": x} for x in metrics if x != "gleu"]  # gleu is unfinished






@parameterized.named_parameters(get_local_metric_names()			)
@for_all_test_methods(
    UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase			)
@local
class   _lowerCamelCase (					parameterized.TestCase			):




       """simple docstring"""
       snake_case      =    {}
       snake_case      =    None



       @pytest.mark.filterwarnings('''ignore:metric_module_factory is deprecated:FutureWarning'''		)
       @pytest.mark.filterwarnings('''ignore:load_metric is deprecated:FutureWarning'''		)
       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE		)->Tuple:




         '''simple docstring'''





         A_      :				List[Any]		     =					'''[...]'''
         A_      :				List[str]		     =					importlib.import_module(
             datasets.load.metric_module_factory(os.path.join('''metrics'''				,       _SCREAMING_SNAKE_CASE		)		).module_path		)
         A_      :				Union[str, Any]		     =					datasets.load.import_main_class(metric_module.__name__				,       dataset=_SCREAMING_SNAKE_CASE		)
         # check parameters
         A_      :				str		     =					inspect.signature(metric._compute		).parameters
         self.assertTrue(all(p.kind != p.VAR_KEYWORD for p in parameters.values()		)		)  # no **kwargs
         # run doctest
         with self.patch_intensive_calls(_SCREAMING_SNAKE_CASE				,       metric_module.__name__		):
           with self.use_local_metrics():
             try:
               A_      :				Union[str, Any]		     =					doctest.testmod(_SCREAMING_SNAKE_CASE				,       verbose=_SCREAMING_SNAKE_CASE				,       raise_on_error=_SCREAMING_SNAKE_CASE		)
             except doctest.UnexpectedException as e:
               raise e.exc_info[1]  # raise the exception that doctest caught
         self.assertEqual(results.failed				,       0		)
         self.assertGreater(results.attempted				,       1		)



       @slow
       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE		)->str:




         '''simple docstring'''





         A_      :				List[Any]		     =					'''[...]'''
         A_      :				Optional[Any]		     =					importlib.import_module(
             datasets.load.metric_module_factory(os.path.join('''metrics'''				,       _SCREAMING_SNAKE_CASE		)		).module_path		)
         # run doctest
         with self.use_local_metrics():
           A_      :				int		     =					doctest.testmod(_SCREAMING_SNAKE_CASE				,       verbose=_SCREAMING_SNAKE_CASE				,       raise_on_error=_SCREAMING_SNAKE_CASE		)
         self.assertEqual(results.failed				,       0		)
         self.assertGreater(results.attempted				,       1		)



       @contextmanager
       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)->Dict:




         '''simple docstring'''





         if metric_name in self.INTENSIVE_CALLS_PATCHER:
           with self.INTENSIVE_CALLS_PATCHER[metric_name](_SCREAMING_SNAKE_CASE		):
             yield
         else:
           yield



       @contextmanager
       def        _snake_case       (      self		)->Optional[int]:




         '''simple docstring'''





         def load_local_metric(_SCREAMING_SNAKE_CASE				,       *_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		):
           return load_metric(os.path.join('''metrics'''				,       _SCREAMING_SNAKE_CASE		)				,       *_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)

         with patch('''datasets.load_metric'''		) as mock_load_metric:
           A_      :				str		     =					load_local_metric
           yield







       @classmethod
       def        _snake_case       (      cls				,       _SCREAMING_SNAKE_CASE		)->List[str]:




         '''simple docstring'''





         def wrapper(_SCREAMING_SNAKE_CASE		):
           A_      :				List[str]		     =					contextmanager(_SCREAMING_SNAKE_CASE		)
           A_      :				int		     =					patcher
           return patcher

         return wrapper




@LocalMetricTest.register_intensive_calls_patcher('''bleurt'''       )
def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  import tensorflow.compat.va as tf
  from bleurt.score import Predictor

  tf.flags.DEFINE_string('''sv''' ,			'''''' ,			''''''       )  # handle pytest cli flags






  class   _lowerCamelCase (					UpperCamelCase			):




         """simple docstring"""



         def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE		)->List[Any]:




           '''simple docstring'''





           assert len(input_dict['''input_ids''']		) == 2
           return np.array([1.0_3, 1.0_4]		)

    # mock predict_fn which is supposed to do a forward pass with a bleurt model
  with patch('''bleurt.score._create_predictor'''       ) as mock_create_predictor:
    A_      :				List[str]		     =					MockedPredictor()
    yield


@LocalMetricTest.register_intensive_calls_patcher('''bertscore'''       )
def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  import torch

  def bert_cos_score_idf(SCREAMING_SNAKE_CASE ,			SCREAMING_SNAKE_CASE ,			*SCREAMING_SNAKE_CASE ,			**SCREAMING_SNAKE_CASE       ):
    return torch.tensor([[1.0, 1.0, 1.0]] * len(SCREAMING_SNAKE_CASE       )       )

  # mock get_model which is supposed to do download a bert model
  # mock bert_cos_score_idf which is supposed to do a forward pass with a bert model
  with patch('''bert_score.scorer.get_model'''       ), patch(
      '''bert_score.scorer.bert_cos_score_idf'''       ) as mock_bert_cos_score_idf:
    A_      :				Optional[int]		     =					bert_cos_score_idf
    yield


@LocalMetricTest.register_intensive_calls_patcher('''comet'''       )
def   _SCREAMING_SNAKE_CASE	( SCREAMING_SNAKE_CASE       ):
  def load_from_checkpoint(SCREAMING_SNAKE_CASE       ):






    class   _lowerCamelCase :




           """simple docstring"""



           def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       *_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE		)->str:




             '''simple docstring'''





             assert len(_SCREAMING_SNAKE_CASE		) == 2
             A_      :				int		     =					[0.1_9, 0.9_2]
             return scores, sum(_SCREAMING_SNAKE_CASE		) / len(_SCREAMING_SNAKE_CASE		)




    return Model()

  # mock load_from_checkpoint which is supposed to do download a bert model
  # mock load_from_checkpoint which is supposed to do download a bert model
  with patch('''comet.download_model'''       ) as mock_download_model:
    A_      :				List[Any]		     =					None
    with patch('''comet.load_from_checkpoint'''       ) as mock_load_from_checkpoint:
      A_      :				Dict		     =					load_from_checkpoint
      yield







def   _SCREAMING_SNAKE_CASE	( ):
  A_      :				Union[str, Any]		     =					load_metric(os.path.join('''metrics''' ,			'''seqeval'''       )       )
  A_      :				Union[str, Any]		     =					'''ERROR'''
  A_      :				Tuple		     =					f'''Scheme should be one of [IOB1, IOB2, IOE1, IOE2, IOBES, BILOU], got {wrong_scheme}'''
  with pytest.raises(SCREAMING_SNAKE_CASE ,			match=re.escape(SCREAMING_SNAKE_CASE       )       ):
    metric.compute(predictions=[] ,			references=[] ,			scheme=SCREAMING_SNAKE_CASE       )








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




       """simple docstring"""



       def __init__(      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=99				,       _SCREAMING_SNAKE_CASE=13				,       _SCREAMING_SNAKE_CASE=7				,       _SCREAMING_SNAKE_CASE=9				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=True				,       _SCREAMING_SNAKE_CASE=False				,       _SCREAMING_SNAKE_CASE=32				,       _SCREAMING_SNAKE_CASE=5				,       _SCREAMING_SNAKE_CASE=4				,       _SCREAMING_SNAKE_CASE=37				,       _SCREAMING_SNAKE_CASE=8				,       _SCREAMING_SNAKE_CASE=0.1				,       _SCREAMING_SNAKE_CASE=0.0_0_2				,       _SCREAMING_SNAKE_CASE=1				,       _SCREAMING_SNAKE_CASE=0				,       _SCREAMING_SNAKE_CASE=0				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       )->str:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					parent
         A_      :				Optional[Any]		     =					batch_size
         A_      :				int		     =					encoder_seq_length
         A_      :				List[Any]		     =					decoder_seq_length
         # For common tests
         A_      :				str		     =					self.decoder_seq_length
         A_      :				int		     =					is_training
         A_      :				Tuple		     =					use_attention_mask
         A_      :				Any		     =					use_labels
         A_      :				List[str]		     =					vocab_size
         A_      :				Dict		     =					hidden_size
         A_      :				Optional[Any]		     =					num_hidden_layers
         A_      :				List[Any]		     =					num_attention_heads
         A_      :				Tuple		     =					d_ff
         A_      :				List[str]		     =					relative_attention_num_buckets
         A_      :				Tuple		     =					dropout_rate
         A_      :				Optional[Any]		     =					initializer_factor
         A_      :				List[str]		     =					eos_token_id
         A_      :				Union[str, Any]		     =					pad_token_id
         A_      :				Any		     =					decoder_start_token_id
         A_      :				int		     =					None
         A_      :				Tuple		     =					decoder_layers



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         return TaConfig.from_pretrained('''google/umt5-base'''		)



       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       _SCREAMING_SNAKE_CASE=None				,       )->List[str]:




         '''simple docstring'''





         if attention_mask is None:
           A_      :				Union[str, Any]		     =					input_ids.ne(config.pad_token_id		)
         if decoder_attention_mask is None:
           A_      :				Any		     =					decoder_input_ids.ne(config.pad_token_id		)
         if head_mask is None:
           A_      :				str		     =					torch.ones(config.num_hidden_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         if decoder_head_mask is None:
           A_      :				List[str]		     =					torch.ones(config.num_decoder_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         if cross_attn_head_mask is None:
           A_      :				Union[str, Any]		     =					torch.ones(
               config.num_decoder_layers				,       config.num_attention_heads				,       device=_SCREAMING_SNAKE_CASE		)
         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        _snake_case       (      self		)->Union[str, Any]:




         '''simple docstring'''





         A_      :				List[Any]		     =					ids_tensor([self.batch_size, self.encoder_seq_length]				,       self.vocab_size		)
         A_      :				int		     =					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_      :				Tuple		     =					input_ids.clamp(self.pad_token_id + 1		)
         A_      :				int		     =					decoder_input_ids.clamp(self.pad_token_id + 1		)

         A_      :				Optional[Any]		     =					self.get_config()
         A_      :				int		     =					config.num_attention_heads
         A_      :				Dict		     =					self.prepare_inputs_dict(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)
         return config, input_dict



       def        _snake_case       (      self		)->Optional[Any]:




         '''simple docstring'''





         A_     ,       A_      :				str		     =					self.prepare_config_and_inputs()
         return config, inputs_dict



       def        _snake_case       (      self		)->Union[str, Any]:




         '''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        _snake_case       (      self		)->Dict:




         '''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        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Dict:




         '''simple docstring'''





         A_      :				Dict		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		)
         model.to(_SCREAMING_SNAKE_CASE		)
         model.eval()
         A_      :				Tuple		     =					model(
             input_ids=_SCREAMING_SNAKE_CASE				,       decoder_input_ids=_SCREAMING_SNAKE_CASE				,       attention_mask=_SCREAMING_SNAKE_CASE				,       decoder_attention_mask=_SCREAMING_SNAKE_CASE				,       )
         A_      :				List[Any]		     =					model(input_ids=_SCREAMING_SNAKE_CASE				,       decoder_input_ids=_SCREAMING_SNAKE_CASE		)
         A_      :				str		     =					result.last_hidden_state
         A_      :				List[Any]		     =					result.past_key_values
         A_      :				Any		     =					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(_SCREAMING_SNAKE_CASE		)				,       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        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Optional[Any]:




         '''simple docstring'''





         A_      :				List[Any]		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		).get_decoder().to(_SCREAMING_SNAKE_CASE		).eval()
         # first forward pass
         A_      :				Union[str, Any]		     =					model(_SCREAMING_SNAKE_CASE				,       use_cache=_SCREAMING_SNAKE_CASE		)
         A_      :				List[str]		     =					model(_SCREAMING_SNAKE_CASE		)
         A_      :				int		     =					model(_SCREAMING_SNAKE_CASE				,       use_cache=_SCREAMING_SNAKE_CASE		)

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

         A_     ,       A_      :				int		     =					outputs.to_tuple()

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

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

         A_      :				Tuple		     =					model(_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']
         A_      :				Union[str, Any]		     =					model(_SCREAMING_SNAKE_CASE				,       past_key_values=_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']

         # select random slice
         A_      :				List[Any]		     =					ids_tensor((1,)				,       output_from_past.shape[-1]		).item()
         A_      :				Union[str, Any]		     =					output_from_no_past[:, -1, random_slice_idx].detach()
         A_      :				List[str]		     =					output_from_past[:, 0, random_slice_idx].detach()

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







       def        _snake_case       (      self				,       _SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE				,       )->Optional[Any]:




         '''simple docstring'''





         A_      :				List[str]		     =					UMTaModel(config=_SCREAMING_SNAKE_CASE		).to(_SCREAMING_SNAKE_CASE		).half().eval()
         A_      :				Tuple		     =					model(**_SCREAMING_SNAKE_CASE		)['''last_hidden_state''']
         self.parent.assertFalse(torch.isnan(_SCREAMING_SNAKE_CASE		).any().item()		)






@require_torch
class   _lowerCamelCase (					UpperCamelCase					,							UpperCamelCase					,							UpperCamelCase					,							unittest.TestCase			):




       """simple docstring"""
       snake_case      =    (
           (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
       )
       snake_case      =    (UMTaForConditionalGeneration,) if is_torch_available() else ()
       snake_case      =    (
           {
               "conversational": UMTaForConditionalGeneration,
               "feature-extraction": UMTaModel,
               "summarization": UMTaForConditionalGeneration,
               "text2text-generation": UMTaForConditionalGeneration,
               "translation": UMTaForConditionalGeneration,
               "question-answering": UMTaForQuestionAnswering,
           }
           if is_torch_available()
           else {}
       )
       snake_case      =    True
       snake_case      =    False
       snake_case      =    False
       snake_case      =    True
       snake_case      =    True
       # The small UMT5 model needs higher percentages for CPU/MP tests
       snake_case      =    [0.8, 0.9]



       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         A_      :				Union[str, Any]		     =					UMTaModelTester(self		)



       @unittest.skip('''Test has a segmentation fault on torch 1.8.0'''		)
       def        _snake_case       (      self		)->List[str]:




         '''simple docstring'''





         A_      :				Any		     =					self.model_tester.prepare_config_and_inputs()
         A_      :				str		     =					UMTaModel(config_and_inputs[0]		).to(_SCREAMING_SNAKE_CASE		)
         with tempfile.TemporaryDirectory() as tmpdirname:
           torch.onnx.export(
               _SCREAMING_SNAKE_CASE				,       (config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])				,       F'''{tmpdirname}/t5_test.onnx'''				,       export_params=_SCREAMING_SNAKE_CASE				,       opset_version=9				,       input_names=['''input_ids''', '''decoder_input_ids''']				,       )



       @unittest.skipIf(torch_device == '''cpu'''				,       '''Cant do half precision'''		)
       def        _snake_case       (      self		)->str:




         '''simple docstring'''





         A_      :				Dict		     =					self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_model_fpaa_forward(*_SCREAMING_SNAKE_CASE		)



       def        _snake_case       (      self		)->int:




         '''simple docstring'''





         A_      :				Optional[int]		     =					['''encoder_attentions''', '''decoder_attentions''', '''cross_attentions''']
         A_      :				List[Any]		     =					self.model_tester.prepare_config_and_inputs()
         A_      :				Union[str, Any]		     =					config_and_inputs[0]
         A_      :				int		     =					UMTaForConditionalGeneration(_SCREAMING_SNAKE_CASE		).eval()
         model.to(_SCREAMING_SNAKE_CASE		)

         A_      :				Any		     =					{
             '''head_mask''': torch.zeros(config.num_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
             '''decoder_head_mask''': torch.zeros(config.num_decoder_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
             '''cross_attn_head_mask''': torch.zeros(config.num_decoder_layers				,       config.num_heads				,       device=_SCREAMING_SNAKE_CASE		),
         }

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

           A_      :				Union[str, Any]		     =					model.generate(
               config_and_inputs[1]['''input_ids''']				,       num_beams=1				,       max_length=3				,       output_attentions=_SCREAMING_SNAKE_CASE				,       return_dict_in_generate=_SCREAMING_SNAKE_CASE				,       **_SCREAMING_SNAKE_CASE				,       )
           # We check the state of decoder_attentions and cross_attentions just from the last step
           A_      :				Tuple		     =					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        _snake_case       (      self		)->Any:




         '''simple docstring'''





         pass






@require_torch
@require_sentencepiece
@require_tokenizers
class   _lowerCamelCase (					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        _snake_case       (      self		)->int:




         '''simple docstring'''





         A_      :				Optional[Any]		     =					UMTaForConditionalGeneration.from_pretrained('''google/umt5-small'''				,       return_dict=_SCREAMING_SNAKE_CASE		).to(_SCREAMING_SNAKE_CASE		)
         A_      :				Union[str, Any]		     =					AutoTokenizer.from_pretrained('''google/umt5-small'''				,       use_fast=_SCREAMING_SNAKE_CASE				,       legacy=_SCREAMING_SNAKE_CASE		)
         A_      :				Tuple		     =					[
             '''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_      :				Union[str, Any]		     =					tokenizer(_SCREAMING_SNAKE_CASE				,       return_tensors='''pt'''				,       padding=_SCREAMING_SNAKE_CASE		).input_ids
         # fmt: off
         A_      :				Dict		     =					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(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)

         A_      :				List[Any]		     =					model.generate(input_ids.to(_SCREAMING_SNAKE_CASE		)		)
         A_      :				str		     =					[
             '''<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_      :				str		     =					tokenizer.batch_decode(_SCREAMING_SNAKE_CASE		)
         self.assertEqual(_SCREAMING_SNAKE_CASE				,       _SCREAMING_SNAKE_CASE		)




