

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



 """simple docstring"""







 def __init__(		self :	Dict     ,						A_ :	Dict     ,						A_ :	Any=9_9     ,						A_ :	str=1_3     ,						A_ :	Union[str, Any]=7     ,						A_ :	Optional[Any]=9     ,						A_ :	Union[str, Any]=True     ,						A_ :	List[str]=True     ,						A_ :	Any=False     ,						A_ :	Dict=3_2     ,						A_ :	Dict=5     ,						A_ :	List[str]=4     ,						A_ :	List[str]=3_7     ,						A_ :	int=8     ,						A_ :	int=0.1     ,						A_ :	int=0.002     ,						A_ :	str=1     ,						A_ :	Any=0     ,						A_ :	Optional[int]=0     ,						A_ :	Any=None     ,						A_ :	Any=None     ,						):






      '''simple docstring'''




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

 def    __magic_name__       (		self :	str       ):






      '''simple docstring'''




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

 def    __magic_name__       (		self :	List[Any]     ,						A_ :	Tuple     ,						A_ :	Optional[int]     ,						A_ :	Tuple     ,						A_ :	Union[str, Any]=None     ,						A_ :	Union[str, Any]=None     ,						A_ :	Union[str, Any]=None     ,						A_ :	str=None     ,						A_ :	int=None     ,						):






      '''simple docstring'''




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






      '''simple docstring'''




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

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

 def    __magic_name__       (		self :	int       ):






      '''simple docstring'''




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

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      return TaConfig(
          vocab_size=1_6_6     ,						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    __magic_name__       (		self :	Tuple       ):






      '''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    __magic_name__       (		self :	Tuple     ,						A_ :	Optional[Any]     ,						A_ :	List[str]     ,						A_ :	Tuple     ,						A_ :	Dict     ,						A_ :	Tuple     ,						A_ :	str     ,						):






      '''simple docstring'''




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






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    UMTaModel(config=A_       ).get_decoder().to(A_       ).eval()
      # first forward pass
      _lowerCAmelCase   :					Tuple        =    model(A_     ,						use_cache=A_       )
      _lowerCAmelCase   :					str        =    model(A_       )
      _lowerCAmelCase   :					Tuple        =    model(A_     ,						use_cache=A_       )

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

      _lowerCAmelCase      ,			_lowerCAmelCase   :					Dict        =    outputs.to_tuple()

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

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

      _lowerCAmelCase   :					List[Any]        =    model(A_       )["last_hidden_state"]
      _lowerCAmelCase   :					Union[str, Any]        =    model(A_     ,						past_key_values=A_       )["last_hidden_state"]

      # select random slice
      _lowerCAmelCase   :					List[str]        =    ids_tensor((1,)     ,						output_from_past.shape[-1]       ).item()
      _lowerCAmelCase   :					Dict        =    output_from_no_past[:, -1, random_slice_idx].detach()
      _lowerCAmelCase   :					Optional[int]        =    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    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	Any     ,						A_ :	Any     ,						):






      '''simple docstring'''




      _lowerCAmelCase   :					Dict        =    UMTaModel(config=A_       ).to(A_       ).half().eval()
      _lowerCAmelCase   :					str        =    model(**A_       )["last_hidden_state"]
      self.parent.assertFalse(torch.isnan(A_       ).any().item()       )





@require_torch
class A__    ( A  ,		A  ,		A  ,		unittest.TestCase ):



 """simple docstring"""







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

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[Any]        =    UMTaModelTester(self       )

 @unittest.skip("Test has a segmentation fault on torch 1.8.0"       )
 def    __magic_name__       (		self :	Any       ):






      '''simple docstring'''




      _lowerCAmelCase   :					List[str]        =    self.model_tester.prepare_config_and_inputs()
      _lowerCAmelCase   :					List[str]        =    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    __magic_name__       (		self :	int       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Tuple        =    self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model_fpaa_forward(*A_       )

 def    __magic_name__       (		self :	List[Any]       ):






      '''simple docstring'''




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

      _lowerCAmelCase   :					Tuple        =    {
          "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()       ):
           _lowerCAmelCase   :					Optional[Any]        =    {name: mask}
           # Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
           if name == "head_mask":
                _lowerCAmelCase   :					Tuple        =    torch.ones(
                    config.num_decoder_layers     ,						config.num_heads     ,						device=A_       )

           _lowerCAmelCase   :					Tuple        =    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
           _lowerCAmelCase   :					Union[str, Any]        =    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    __magic_name__       (		self :	Optional[int]       ):






      '''simple docstring'''




      pass





@require_torch
@require_sentencepiece
@require_tokenizers
class A__    ( 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    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      _lowerCAmelCase   :					Optional[int]        =    UMTaForConditionalGeneration.from_pretrained("google/umt5-small"     ,						return_dict=A_       ).to(A_       )
      _lowerCAmelCase   :					Optional[int]        =    AutoTokenizer.from_pretrained("google/umt5-small"     ,						use_fast=A_     ,						legacy=A_       )
      _lowerCAmelCase   :					Any        =    [
          "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>.",
      ]
      _lowerCAmelCase   :					Optional[Any]        =    tokenizer(A_     ,						return_tensors="pt"     ,						padding=A_       ).input_ids
      # fmt: off
      _lowerCAmelCase   :					Any        =    torch.tensor(
          [
              [ 3_8_5_3_0, 2_1_0_7_0_3, 2_5_6_2_9_9, 1_4_1_0, 2_5_6_2_9_8, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
              [   8_2_6, 3_2_1, 6_7_1, 2_5_9_2_2, 2_5_6_2_9_9, 2_7_4, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
              [  1_4_6_0, 3_3_9, 3_1_2, 1_9_0_1_4, 1_0_6_2_0, 7_5_8, 2_5_6_2_9_9, 2_3_5_5,2_7_4, 1, 0, 0, 0, 0, 0, 0,0, 0],
              [   5_1_7, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 3_0_1, 2_5_6_2_9_8, 2_7_5, 1_1_9_9_8_3,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
              [   3_2_0, 2_5_6_2_9_9, 1_4_8_6_9, 2_8_1, 2_2_3_4, 2_8_9, 2_2_7_5, 3_3_3,6_1_3_9_1, 2_8_9, 2_5_6_2_9_8, 5_4_3, 2_5_6_2_9_7, 1_6_8_7_1_4, 3_2_9, 2_5_6_2_9_6,2_7_4, 1],
          ]       )
      # fmt: on
      torch.testing.assert_allclose(A_     ,						A_       )

      _lowerCAmelCase   :					List[Any]        =    model.generate(input_ids.to(A_       )       )
      _lowerCAmelCase   :					List[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>",
      ]
      _lowerCAmelCase   :					Tuple        =    tokenizer.batch_decode(A_       )
      self.assertEqual(A_     ,						A_       )









from pathlib import PurePosixPath
from typing import Optional

import fsspec
from fsspec import AbstractFileSystem
from huggingface_hub.hf_api import DatasetInfo

from ..utils.file_utils import get_authentication_headers_for_url
from ..utils.hub import hf_hub_url





class A__    ( A ):



 """simple docstring"""







 _lowercase     :					List[str]      =						''''''
 _lowercase     :					Optional[Any]      =						'''hf-legacy'''  # "hf://"" is reserved for hffs

 def __init__(		self :	Dict     ,						A_ :	Optional[DatasetInfo] = None     ,						A_ :	Optional[str] = None     ,						**A_ :	int     ,						):






      '''simple docstring'''




      super().__init__(self     ,						**A_       )
      _lowerCAmelCase   :					Union[str, Any]        =    repo_info
      _lowerCAmelCase   :					Optional[int]        =    token
      _lowerCAmelCase   :					str        =    None

 def    __magic_name__       (		self :	Tuple       ):






      '''simple docstring'''




      if self.dir_cache is None:
           _lowerCAmelCase   :					int        =    {}
           for hf_file in self.repo_info.siblings:
                # TODO(QL): add sizes
                _lowerCAmelCase   :					Optional[int]        =    {
                    "name": hf_file.rfilename,
                    "size": None,
                    "type": "file",
                }
                self.dir_cache.update(
                    {
                        str(A_       ): {"name": str(A_       ), "size": None, "type": "directory"}
                        for d in list(PurePosixPath(hf_file.rfilename       ).parents       )[:-1]
                    }       )

 def    __magic_name__       (		self :	Dict     ,						A_ :	str     ,						A_ :	str = "rb"     ,						**A_ :	Optional[int]     ,						):






      '''simple docstring'''




      if not isinstance(self.repo_info     ,						A_       ):
           raise NotImplementedError(F'''Open is only implemented for dataset repositories, but got {self.repo_info}'''       )
      _lowerCAmelCase   :					str        =    hf_hub_url(self.repo_info.id     ,						A_     ,						revision=self.repo_info.sha       )
      return fsspec.open(
          A_     ,						mode=A_     ,						headers=get_authentication_headers_for_url(A_     ,						use_auth_token=self.token       )     ,						client_kwargs={"trust_env": True}     ,						).open()

 def    __magic_name__       (		self :	Union[str, Any]     ,						A_ :	List[Any]     ,						**A_ :	str       ):






      '''simple docstring'''




      self._get_dirs()
      _lowerCAmelCase   :					Optional[int]        =    self._strip_protocol(A_       )
      if path in self.dir_cache:
           return self.dir_cache[path]
      else:
           raise FileNotFoundError(A_       )


 def    __magic_name__       (		self :	Any     ,						A_ :	str     ,						A_ :	Any=False     ,						**A_ :	Dict       ):






      '''simple docstring'''




      self._get_dirs()
      _lowerCAmelCase   :					Dict        =    PurePosixPath(path.strip("/"       )       )
      _lowerCAmelCase   :					str        =    {}
      for p, f in self.dir_cache.items():
           _lowerCAmelCase   :					List[str]        =    PurePosixPath(p.strip("/"       )       )
           _lowerCAmelCase   :					Tuple        =    p.parent
           if root == path:
                _lowerCAmelCase   :					Union[str, Any]        =    f
      _lowerCAmelCase   :					List[str]        =    list(paths.values()       )
      if detail:
           return out
      else:
           return sorted(f["name"] for f in out       )






