





import unittest

from transformers import SpeechaTextaConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

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


if is_torch_available():
     import torch

     from transformers.models.speech_to_text_a.modeling_speech_to_text_a import (
         SpeechaTextaDecoder,
         SpeechaTextaForCausalLM,
     )



@require_torch
class 				_UpperCamelCase       :

       def __init__(		self:			Any       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			int=99       , _SCREAMING_SNAKE_CASE:			Optional[Any]=13       , _SCREAMING_SNAKE_CASE:			int=16       , _SCREAMING_SNAKE_CASE:			str=7       , _SCREAMING_SNAKE_CASE:			Optional[int]=True       , _SCREAMING_SNAKE_CASE:			Tuple=True       , _SCREAMING_SNAKE_CASE:			List[Any]=True       , _SCREAMING_SNAKE_CASE:			Optional[int]=False       , _SCREAMING_SNAKE_CASE:			List[str]=True       , _SCREAMING_SNAKE_CASE:			Any=2       , _SCREAMING_SNAKE_CASE:			Tuple=32       , _SCREAMING_SNAKE_CASE:			str=4       , _SCREAMING_SNAKE_CASE:			List[Any]=4       , _SCREAMING_SNAKE_CASE:			Any=30       , _SCREAMING_SNAKE_CASE:			Optional[int]=0       , _SCREAMING_SNAKE_CASE:			List[Any]=1       , _SCREAMING_SNAKE_CASE:			Optional[Any]=2       , _SCREAMING_SNAKE_CASE:			Optional[Any]=None       , )							->			List[str]:




              """simple docstring"""


              UpperCamelCase_ 							=	parent
              UpperCamelCase_ 							=	batch_size
              UpperCamelCase_ 							=	decoder_seq_length
              # For common tests
              UpperCamelCase_ 							=	self.decoder_seq_length
              UpperCamelCase_ 							=	is_training
              UpperCamelCase_ 							=	use_attention_mask
              UpperCamelCase_ 							=	use_labels

              UpperCamelCase_ 							=	vocab_size
              UpperCamelCase_ 							=	d_model
              UpperCamelCase_ 							=	d_model
              UpperCamelCase_ 							=	decoder_layers
              UpperCamelCase_ 							=	decoder_layers
              UpperCamelCase_ 							=	decoder_ffn_dim
              UpperCamelCase_ 							=	decoder_attention_heads
              UpperCamelCase_ 							=	decoder_attention_heads
              UpperCamelCase_ 							=	eos_token_id
              UpperCamelCase_ 							=	bos_token_id
              UpperCamelCase_ 							=	pad_token_id
              UpperCamelCase_ 							=	decoder_start_token_id
              UpperCamelCase_ 							=	use_cache
              UpperCamelCase_ 							=	max_position_embeddings

              UpperCamelCase_ 							=	None
              UpperCamelCase_ 							=	decoder_seq_length
              UpperCamelCase_ 							=	2
              UpperCamelCase_ 							=	1

       def lowercase			(		self:			Dict		)							->			List[str]:




              """simple docstring"""


              UpperCamelCase_ 							=	ids_tensor([self.batch_size, self.decoder_seq_length]       , self.vocab_size		)

              UpperCamelCase_ 							=	None
              if self.use_attention_mask:
                     UpperCamelCase_ 							=	ids_tensor([self.batch_size, self.decoder_seq_length]       , vocab_size=2		)

              UpperCamelCase_ 							=	None
              if self.use_labels:
                     UpperCamelCase_ 							=	ids_tensor([self.batch_size, self.decoder_seq_length]       , self.vocab_size		)

              UpperCamelCase_ 							=	SpeechaTextaConfig(
                  vocab_size=self.vocab_size       , d_model=self.d_model       , decoder_layers=self.decoder_layers       , decoder_ffn_dim=self.decoder_ffn_dim       , decoder_attention_heads=self.decoder_attention_heads       , eos_token_id=self.eos_token_id       , bos_token_id=self.bos_token_id       , use_cache=self.use_cache       , pad_token_id=self.pad_token_id       , decoder_start_token_id=self.decoder_start_token_id       , max_position_embeddings=self.max_position_embeddings       , )

              return (
                  config,
                  input_ids,
                  attention_mask,
                  lm_labels,
              )

       def lowercase			(		self:			Tuple       , _SCREAMING_SNAKE_CASE:			Optional[Any]       , _SCREAMING_SNAKE_CASE:			Optional[int]       , _SCREAMING_SNAKE_CASE:			Dict       , _SCREAMING_SNAKE_CASE:			Optional[int]       , )							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	True
              UpperCamelCase_ 							=	SpeechaTextaDecoder(config=__UpperCamelCase		).to(__UpperCamelCase		).eval()
              UpperCamelCase_ 							=	input_ids[:2]

              input_ids[input_ids == 0] += 1
              # first forward pass
              UpperCamelCase_ 							=	model(__UpperCamelCase       , use_cache=__UpperCamelCase		)
              UpperCamelCase_ 							=	model(__UpperCamelCase		)
              UpperCamelCase_ 							=	model(__UpperCamelCase       , use_cache=__UpperCamelCase		)

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

              UpperCamelCase_ 							=	outputs["past_key_values"]

              # create hypothetical next token and extent to next_input_ids
              UpperCamelCase_ 							=	ids_tensor((2, 1)       , config.vocab_size - 1		) + 1

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

              print(__UpperCamelCase		)

              UpperCamelCase_ 							=	model(__UpperCamelCase		)["last_hidden_state"]
              UpperCamelCase_ 							=	model(__UpperCamelCase       , past_key_values=__UpperCamelCase		)["last_hidden_state"]

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

              # test that outputs are equal for slice
              assert torch.allclose(__UpperCamelCase       , __UpperCamelCase       , atol=1e-3		)



       def lowercase			(		self:			str		)							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	self.prepare_config_and_inputs()
              (
                  (
                  UpperCamelCase_      
              )							,			(
                  UpperCamelCase_      
              )							,			(
                  UpperCamelCase_      
              )							,			(
                  UpperCamelCase_      
              )							,			      
              ) 							=	config_and_inputs

              UpperCamelCase_ 							=	{
                  "input_ids": input_ids,
                  "attention_mask": attention_mask,
              }
              return config, inputs_dict



@require_torch
class 				_UpperCamelCase       (						lowercase_	,    lowercase_	,    lowercase_	,    unittest.TestCase  ):
       _UpperCamelCase					:       List[str]       = (SpeechaTextaDecoder, SpeechaTextaForCausalLM) if is_torch_available() else ()
       _UpperCamelCase					:       List[Any]       = (SpeechaTextaForCausalLM,) if is_torch_available() else ()
       _UpperCamelCase					:       List[str]       = {'''text-generation''': SpeechaTextaForCausalLM} if is_torch_available() else {}
       _UpperCamelCase					:       Dict       = True
       _UpperCamelCase					:       Optional[int]       = False

       def lowercase			(		self:			str       , )							->			Any:




              """simple docstring"""


              UpperCamelCase_ 							=	SpeechaTextaStandaloneDecoderModelTester(self       , is_training=__UpperCamelCase		)
              UpperCamelCase_ 							=	ConfigTester(self       , config_class=__UpperCamelCase		)

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




              """simple docstring"""


              pass

       def lowercase			(		self:			Tuple		)							->			Any:




              """simple docstring"""


              pass

       def lowercase			(		self:			Optional[Any]		)							->			List[str]:




              """simple docstring"""


              pass

       def lowercase			(		self:			Tuple		)							->			Tuple:




              """simple docstring"""


              self.config_tester.run_common_tests()

       def lowercase			(		self:			List[Any]		)							->			Tuple:




              """simple docstring"""


              UpperCamelCase_ 							=	self.model_tester.prepare_config_and_inputs()
              self.model_tester.create_and_check_decoder_model_past(*__UpperCamelCase		)



       def lowercase			(		self:			List[Any]		)							->			List[Any]:




              """simple docstring"""


              return








"""simple docstring"""
import os
from shutil import copyfile
from typing import List, Optional, Tuple

from tokenizers import processors

from ...tokenization_utils import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
		from .tokenization_mbartaa import MBartaaTokenizer
else:
		_snake_case	:						Dict					=					None


_snake_case	:						int					=					logging.get_logger(__name__)

_snake_case	:						Optional[Any]					=					{'vocab_file': 'sentencepiece.bpe.model', 'tokenizer_file': 'tokenizer.json'}

_snake_case	:						str					=					{
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    },
    'tokenizer_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/tokenizer.json'
        ),
    },
}

_snake_case	:						Union[str, Any]					=					{
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
_snake_case	:						int					=					['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']






class    _UpperCAmelCase  (      lowercase_					):
					UpperCamelCase						=				VOCAB_FILES_NAMES
					UpperCamelCase						=				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					UpperCamelCase						=				PRETRAINED_VOCAB_FILES_MAP
					UpperCamelCase						=				['''input_ids''', '''attention_mask''']
					UpperCamelCase						=				MBartaaTokenizer

					UpperCamelCase						=				[]
					UpperCamelCase						=				[]


					def __init__(   self     :Any	,			__UpperCamelCase     :str=None	,			__UpperCamelCase     :Any=None	,			__UpperCamelCase     :int=None	,			__UpperCamelCase     :Union[str, Any]=None	,			__UpperCamelCase     :Optional[Any]="</s>"	,			__UpperCamelCase     :Tuple="</s>"	,			__UpperCamelCase     :Any="<s>"	,			__UpperCamelCase     :Optional[int]="<unk>"	,			__UpperCamelCase     :str="<pad>"	,			__UpperCamelCase     :int="<mask>"	,			**__UpperCamelCase     :List[str]	,			):
												# Mask token behave like a normal word, i.e. include the space before it
												A						=    AddedToken(__UpperCamelCase	,			lstrip=__UpperCamelCase	,			rstrip=__UpperCamelCase						) if isinstance(__UpperCamelCase	,			__UpperCamelCase						) else mask_token

												A						=    kwargs.get("additional_special_tokens"	,			[]						)
												kwargs["additional_special_tokens"] += [
												    code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
												]

												super().__init__(
												    __UpperCamelCase	,			src_lang=__UpperCamelCase	,			tgt_lang=__UpperCamelCase	,			tokenizer_file=__UpperCamelCase	,			eos_token=__UpperCamelCase	,			sep_token=__UpperCamelCase	,			cls_token=__UpperCamelCase	,			unk_token=__UpperCamelCase	,			pad_token=__UpperCamelCase	,			mask_token=__UpperCamelCase	,			**__UpperCamelCase	,			)

												A						=    vocab_file
												A						=    False if not self.vocab_file else True

												A						=    {
												    lang_code: self.convert_tokens_to_ids(__UpperCamelCase						) for lang_code in FAIRSEQ_LANGUAGE_CODES
												}

												A						=    src_lang if src_lang is not None else "en_XX"
												A						=    tgt_lang
												A						=    self.lang_code_to_id[self._src_lang]
												self.set_src_lang_special_tokens(self._src_lang						)


					@property
					def       lowerCamelCase			(   self     :int						):
												return self._src_lang


					@src_lang.setter
					def       lowerCamelCase			(   self     :Optional[int]	,			__UpperCamelCase     :str						):
												A						=    new_src_lang
												self.set_src_lang_special_tokens(self._src_lang						)


					def       lowerCamelCase			(   self     :Any	,			__UpperCamelCase     :List[int]	,			__UpperCamelCase     :Optional[List[int]] = None						):
												if token_ids_a is None:
																			return self.prefix_tokens + token_ids_a + self.suffix_tokens
												# We don't expect to process pairs, but leave the pair logic for API consistency
												return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens


					def       lowerCamelCase			(   self     :Optional[Any]	,			__UpperCamelCase     :List[str]	,			__UpperCamelCase     :str = "en_XX"	,			__UpperCamelCase     :Optional[List[str]] = None	,			__UpperCamelCase     :str = "ro_RO"	,			**__UpperCamelCase     :Tuple	,			):
												A						=    src_lang
												A						=    tgt_lang
												return super().prepare_seqaseq_batch(__UpperCamelCase	,			__UpperCamelCase	,			**__UpperCamelCase						)


					def       lowerCamelCase			(   self     :str						):
												return self.set_src_lang_special_tokens(self.src_lang						)


					def       lowerCamelCase			(   self     :Optional[Any]						):
												return self.set_tgt_lang_special_tokens(self.tgt_lang						)


					def       lowerCamelCase			(   self     :Any	,			__UpperCamelCase     :str						):
												A						=    self.convert_tokens_to_ids(__UpperCamelCase						)
												A						=    [self.cur_lang_code_id]
												A						=    [self.eos_token_id]

												A						=    self.convert_ids_to_tokens(self.prefix_tokens						)
												A						=    self.convert_ids_to_tokens(self.suffix_tokens						)

												A						=    processors.TemplateProcessing(
												    single=prefix_tokens_str + ["$A"] + suffix_tokens_str	,			pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str	,			special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str	,			self.prefix_tokens + self.suffix_tokens						)						)	,			)


					def       lowerCamelCase			(   self     :Optional[Any]	,			__UpperCamelCase     :str						):
												A						=    self.convert_tokens_to_ids(__UpperCamelCase						)
												A						=    [self.cur_lang_code_id]
												A						=    [self.eos_token_id]

												A						=    self.convert_ids_to_tokens(self.prefix_tokens						)
												A						=    self.convert_ids_to_tokens(self.suffix_tokens						)

												A						=    processors.TemplateProcessing(
												    single=prefix_tokens_str + ["$A"] + suffix_tokens_str	,			pair=prefix_tokens_str + ["$A", "$B"] + suffix_tokens_str	,			special_tokens=list(zip(prefix_tokens_str + suffix_tokens_str	,			self.prefix_tokens + self.suffix_tokens						)						)	,			)


					def       lowerCamelCase			(   self     :Optional[Any]	,			__UpperCamelCase     :Any	,			__UpperCamelCase     :str	,			__UpperCamelCase     :Optional[str]	,			__UpperCamelCase     :Optional[str]	,			**__UpperCamelCase     :Union[str, Any]						):
												if src_lang is None or tgt_lang is None:
																			raise ValueError("Translation requires a `src_lang` and a `tgt_lang` for this model"						)
												A						=    src_lang
												A						=    self(__UpperCamelCase	,			add_special_tokens=__UpperCamelCase	,			return_tensors=__UpperCamelCase	,			**__UpperCamelCase						)
												A						=    self.convert_tokens_to_ids(__UpperCamelCase						)
												A						=    tgt_lang_id
												return inputs


					def       lowerCamelCase			(   self     :Tuple	,			__UpperCamelCase     :str	,			__UpperCamelCase     :Optional[str] = None						):
												if not self.can_save_slow_tokenizer:
																			raise ValueError(
																			    "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
																			    "tokenizer."						)

												if not os.path.isdir(__UpperCamelCase						):
																			logger.error(f"Vocabulary path ({save_directory}) should be a directory"						)
																			return
												A						=    os.path.join(
												    __UpperCamelCase	,			(filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]						)

												if os.path.abspath(self.vocab_file						) != os.path.abspath(__UpperCamelCase						):
																			copyfile(self.vocab_file	,			__UpperCamelCase						)

												return (out_vocab_file,)


