


'''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:
						_A      =     None


_A      =     logging.get_logger(__name__)

_A      =     {"""vocab_file""": """sentencepiece.bpe.model""", """tokenizer_file""": """tokenizer.json"""}

_A      =     {
    """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"""
        ),
    },
}

_A      =     {
    """facebook/mbart-large-50-one-to-many-mmt""": 10_24,
}

# fmt: off
_A      =     ["""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      SCREAMING_SNAKE_CASE_	( snake_case		):
			__a					:      int					=  VOCAB_FILES_NAMES
			__a					:      str					=  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			__a					:      List[str]					=  PRETRAINED_VOCAB_FILES_MAP
			__a					:      List[Any]					=  ['''input_ids''', '''attention_mask''']
			__a					:      str					=  MBartaaTokenizer

			__a					:      List[int]					=  []
			__a					:      List[int]					=  []




			def __init__(      self					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase=None					,      lowercase="</s>"					,      lowercase="</s>"					,      lowercase="<s>"					,      lowercase="<unk>"					,      lowercase="<pad>"					,      lowercase="<mask>"					,      **lowercase					,      )			->						Optional[Any]:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = AddedToken(lowercase					,      lstrip=lowercase					,      rstrip=lowercase       ) if isinstance(lowercase					,      lowercase       ) else mask_token

								__SCREAMING_SNAKE_CASE					:							List[Any]			   = 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__(
								    lowercase					,      src_lang=lowercase					,      tgt_lang=lowercase					,      tokenizer_file=lowercase					,      eos_token=lowercase					,      sep_token=lowercase					,      cls_token=lowercase					,      unk_token=lowercase					,      pad_token=lowercase					,      mask_token=lowercase					,      **lowercase					,      )

								__SCREAMING_SNAKE_CASE					:							Any			   = vocab_file
								__SCREAMING_SNAKE_CASE					:							int			   = False if not self.vocab_file else True

								__SCREAMING_SNAKE_CASE					:							Tuple			   = {
								    lang_code: self.convert_tokens_to_ids(lowercase       ) for lang_code in FAIRSEQ_LANGUAGE_CODES
								}

								__SCREAMING_SNAKE_CASE					:							Any			   = src_lang if src_lang is not None else '''en_XX'''
								__SCREAMING_SNAKE_CASE					:							Any			   = tgt_lang
								__SCREAMING_SNAKE_CASE					:							Dict			   = self.lang_code_to_id[self._src_lang]
								self.set_src_lang_special_tokens(self._src_lang       )




			@property
			def 			_snake_case   (      self       )			->						str:


								'''simple docstring'''






								return self._src_lang




			@src_lang.setter
			def 			_snake_case   (      self					,      lowercase       )			->						None:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							List[Any]			   = new_src_lang
								self.set_src_lang_special_tokens(self._src_lang       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase = None       )			->						List[int]:


								'''simple docstring'''






								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 			_snake_case   (      self					,      lowercase					,      lowercase = "en_XX"					,      lowercase = None					,      lowercase = "ro_RO"					,      **lowercase					,      )			->						BatchEncoding:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = src_lang
								__SCREAMING_SNAKE_CASE					:							str			   = tgt_lang
								return super().prepare_seqaseq_batch(lowercase					,      lowercase					,      **lowercase       )




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


								'''simple docstring'''






								return self.set_src_lang_special_tokens(self.src_lang       )




			def 			_snake_case   (      self       )			->						str:


								'''simple docstring'''






								return self.set_tgt_lang_special_tokens(self.tgt_lang       )




			def 			_snake_case   (      self					,      lowercase       )			->						None:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Dict			   = self.convert_tokens_to_ids(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = [self.cur_lang_code_id]
								__SCREAMING_SNAKE_CASE					:							List[str]			   = [self.eos_token_id]

								__SCREAMING_SNAKE_CASE					:							int			   = self.convert_ids_to_tokens(self.prefix_tokens       )
								__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = self.convert_ids_to_tokens(self.suffix_tokens       )

								__SCREAMING_SNAKE_CASE					:							str			   = 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 			_snake_case   (      self					,      lowercase       )			->						None:


								'''simple docstring'''






								__SCREAMING_SNAKE_CASE					:							Any			   = self.convert_tokens_to_ids(lowercase       )
								__SCREAMING_SNAKE_CASE					:							Any			   = [self.cur_lang_code_id]
								__SCREAMING_SNAKE_CASE					:							List[str]			   = [self.eos_token_id]

								__SCREAMING_SNAKE_CASE					:							Any			   = self.convert_ids_to_tokens(self.prefix_tokens       )
								__SCREAMING_SNAKE_CASE					:							Dict			   = self.convert_ids_to_tokens(self.suffix_tokens       )

								__SCREAMING_SNAKE_CASE					:							str			   = 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 			_snake_case   (      self					,      lowercase					,      lowercase					,      lowercase					,      lowercase					,      **lowercase       )			->						Union[str, Any]:


								'''simple docstring'''






								if src_lang is None or tgt_lang is None:
													raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''       )
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = src_lang
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = self(lowercase					,      add_special_tokens=lowercase					,      return_tensors=lowercase					,      **lowercase       )
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.convert_tokens_to_ids(lowercase       )
								__SCREAMING_SNAKE_CASE					:							int			   = tgt_lang_id
								return inputs




			def 			_snake_case   (      self					,      lowercase					,      lowercase = None       )			->						Tuple[str]:


								'''simple docstring'''






								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(lowercase       ):
													logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""       )
													return
								__SCREAMING_SNAKE_CASE					:							str			   = os.path.join(
								    lowercase					,      (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']       )

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

								return (out_vocab_file,)






'''simple docstring'''






from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin






@dataclass
class      SCREAMING_SNAKE_CASE_	( snake_case		):
			__a					:      torch.FloatTensor






class      SCREAMING_SNAKE_CASE_	( snake_case					,					snake_case		):




			@register_to_config
			def __init__(      self					,      lowercase = 1_6					,      lowercase = 8_8					,      lowercase = None					,      lowercase = None					,      lowercase = 1					,      lowercase = 0.0					,      lowercase = 3_2					,      lowercase = None					,      lowercase = False					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = "geglu"					,      lowercase = None					,      lowercase = False					,      lowercase = False					,      lowercase = False					,      lowercase = "layer_norm"					,      lowercase = True					,      )			->						Optional[int]:


								'''simple docstring'''






								super().__init__()
								__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = use_linear_projection
								__SCREAMING_SNAKE_CASE					:							Dict			   = num_attention_heads
								__SCREAMING_SNAKE_CASE					:							Optional[int]			   = attention_head_dim
								__SCREAMING_SNAKE_CASE					:							Any			   = num_attention_heads * attention_head_dim

								# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
								# Define whether input is continuous or discrete depending on configuration
								__SCREAMING_SNAKE_CASE					:							List[Any]			   = (in_channels is not None) and (patch_size is None)
								__SCREAMING_SNAKE_CASE					:							List[str]			   = num_vector_embeds is not None
								__SCREAMING_SNAKE_CASE					:							Any			   = in_channels is not None and patch_size is not None

								if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
													__SCREAMING_SNAKE_CASE					:							Any			   = (
													    f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
													    ''' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'''
													    ''' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'''
													    ''' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'''
													    ''' would be very nice if you could open a Pull request for the `transformer/config.json` file'''
													)
													deprecate('''norm_type!=num_embeds_ada_norm'''					,      '''1.0.0'''					,      lowercase					,      standard_warn=lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = '''ada_norm'''

								if self.is_input_continuous and self.is_input_vectorized:
													raise ValueError(
													    f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
													    ''' sure that either `in_channels` or `num_vector_embeds` is None.'''       )
								elif self.is_input_vectorized and self.is_input_patches:
													raise ValueError(
													    f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
													    ''' sure that either `num_vector_embeds` or `num_patches` is None.'''       )
								elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
													raise ValueError(
													    f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
													    f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."""       )

								# 2. Define input layers
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = in_channels

													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = torch.nn.GroupNorm(num_groups=lowercase					,      num_channels=lowercase					,      eps=1e-6					,      affine=lowercase       )
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							List[Any]			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Any			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
													assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = sample_size
													__SCREAMING_SNAKE_CASE					:							str			   = sample_size
													__SCREAMING_SNAKE_CASE					:							int			   = num_vector_embeds
													__SCREAMING_SNAKE_CASE					:							str			   = self.height * self.width

													__SCREAMING_SNAKE_CASE					:							int			   = ImagePositionalEmbeddings(
													    num_embed=lowercase					,      embed_dim=lowercase					,      height=self.height					,      width=self.width       )
								elif self.is_input_patches:
													assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size
													__SCREAMING_SNAKE_CASE					:							Any			   = sample_size

													__SCREAMING_SNAKE_CASE					:							int			   = patch_size
													__SCREAMING_SNAKE_CASE					:							Any			   = PatchEmbed(
													    height=lowercase					,      width=lowercase					,      patch_size=lowercase					,      in_channels=lowercase					,      embed_dim=lowercase					,      )

								# 3. Define transformers blocks
								__SCREAMING_SNAKE_CASE					:							Tuple			   = nn.ModuleList(
								    [
								        BasicTransformerBlock(
								            lowercase					,      lowercase					,      lowercase					,      dropout=lowercase					,      cross_attention_dim=lowercase					,      activation_fn=lowercase					,      num_embeds_ada_norm=lowercase					,      attention_bias=lowercase					,      only_cross_attention=lowercase					,      upcast_attention=lowercase					,      norm_type=lowercase					,      norm_elementwise_affine=lowercase					,      )
								        for d in range(lowercase       )
								    ]       )

								# 4. Define output layers
								__SCREAMING_SNAKE_CASE					:							List[str]			   = in_channels if out_channels is None else out_channels
								if self.is_input_continuous:
													# TODO: should use out_channels for continuous projections
													if use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							str			   = nn.Linear(lowercase					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.Convad(lowercase					,      lowercase					,      kernel_size=1					,      stride=1					,      padding=0       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = nn.LayerNorm(lowercase       )
													__SCREAMING_SNAKE_CASE					:							int			   = nn.Linear(lowercase					,      self.num_vector_embeds - 1       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = nn.LayerNorm(lowercase					,      elementwise_affine=lowercase					,      eps=1e-6       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      2 * inner_dim       )
													__SCREAMING_SNAKE_CASE					:							Any			   = nn.Linear(lowercase					,      patch_size * patch_size * self.out_channels       )




			def 			_snake_case   (      self					,      lowercase					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = None					,      lowercase = True					,      )			->						Optional[Any]:


								'''simple docstring'''






								if attention_mask is not None and attention_mask.ndim == 2:
													# assume that mask is expressed as:
													#   (1 = keep,      0 = discard)
													# convert mask into a bias that can be added to attention scores:
													#       (keep = +0,     discard = -10000.0)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = (1 - attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							Dict			   = attention_mask.unsqueeze(1       )

								# convert encoder_attention_mask to a bias the same way we do for attention_mask
								if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
													__SCREAMING_SNAKE_CASE					:							Optional[Any]			   = (1 - encoder_attention_mask.to(hidden_states.dtype       )) * -1_0_0_0_0.0
													__SCREAMING_SNAKE_CASE					:							List[Any]			   = encoder_attention_mask.unsqueeze(1       )

								# 1. Input
								if self.is_input_continuous:
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[Any]			   = hidden_states.shape
													__SCREAMING_SNAKE_CASE					:							Dict			   = hidden_states

													__SCREAMING_SNAKE_CASE					:							str			   = self.norm(lowercase       )
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							int			   = self.proj_in(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							List[str]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.shape[1]
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = hidden_states.permute(0					,      2					,      3					,      1       ).reshape(lowercase					,      height * width					,      lowercase       )
																		__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_in(lowercase       )
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.latent_image_embedding(lowercase       )
								elif self.is_input_patches:
													__SCREAMING_SNAKE_CASE					:							int			   = self.pos_embed(lowercase       )

								# 2. Blocks
								for block in self.transformer_blocks:
													__SCREAMING_SNAKE_CASE					:							Tuple			   = block(
													    lowercase					,      attention_mask=lowercase					,      encoder_hidden_states=lowercase					,      encoder_attention_mask=lowercase					,      timestep=lowercase					,      cross_attention_kwargs=lowercase					,      class_labels=lowercase					,      )

								# 3. Output
								if self.is_input_continuous:
													if not self.use_linear_projection:
																		__SCREAMING_SNAKE_CASE					:							Tuple			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()
																		__SCREAMING_SNAKE_CASE					:							Union[str, Any]			   = self.proj_out(lowercase       )
													else:
																		__SCREAMING_SNAKE_CASE					:							str			   = self.proj_out(lowercase       )
																		__SCREAMING_SNAKE_CASE					:							str			   = hidden_states.reshape(lowercase					,      lowercase					,      lowercase					,      lowercase       ).permute(0					,      3					,      1					,      2       ).contiguous()

													__SCREAMING_SNAKE_CASE					:							str			   = hidden_states + residual
								elif self.is_input_vectorized:
													__SCREAMING_SNAKE_CASE					:							Dict			   = self.norm_out(lowercase       )
													__SCREAMING_SNAKE_CASE					:							str			   = self.out(lowercase       )
													# (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = logits.permute(0					,      2					,      1       )

													# log(p(x_0))
													__SCREAMING_SNAKE_CASE					:							int			   = F.log_softmax(logits.double()					,      dim=1       ).float()
								elif self.is_input_patches:
													# TODO: cleanup!
													__SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.transformer_blocks[0].norma.emb(
													    lowercase					,      lowercase					,      hidden_dtype=hidden_states.dtype       )
													__SCREAMING_SNAKE_CASE   ,   __SCREAMING_SNAKE_CASE					:							Optional[int]			   = self.proj_out_a(F.silu(lowercase       )       ).chunk(2					,      dim=1       )
													__SCREAMING_SNAKE_CASE					:							int			   = self.norm_out(lowercase       ) * (1 + scale[:, None]) + shift[:, None]
													__SCREAMING_SNAKE_CASE					:							int			   = self.proj_out_a(lowercase       )

													# unpatchify
													__SCREAMING_SNAKE_CASE					:							List[str]			   = int(hidden_states.shape[1] ** 0.5       )
													__SCREAMING_SNAKE_CASE					:							int			   = hidden_states.reshape(
													    shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels)       )
													__SCREAMING_SNAKE_CASE					:							int			   = torch.einsum('''nhwpqc->nchpwq'''					,      lowercase       )
													__SCREAMING_SNAKE_CASE					:							Any			   = hidden_states.reshape(
													    shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size)       )

								if not return_dict:
													return (output,)

								return TransformeraDModelOutput(sample=lowercase       )


