




"""simple docstring"""

import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


__magic_name__				     =     logging.get_logger(__name__)

__magic_name__				     =     """▁"""

__magic_name__				     =     {"""vocab_file""": """sentencepiece.bpe.model"""}

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

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

# fmt: off
__magic_name__				     =     ["""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__   (						_SCREAMING_SNAKE_CASE      ):
   snake_case			   =       VOCAB_FILES_NAMES
   snake_case			   =       PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   snake_case			   =       PRETRAINED_VOCAB_FILES_MAP
   snake_case			   =       ["""input_ids""", """attention_mask"""]

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

          lowerCamelCase__          =						{} if sp_model_kwargs is None else sp_model_kwargs

          lowerCamelCase__          =						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__(
              src_lang=SCREAMING_SNAKE_CASE_   ,	tgt_lang=SCREAMING_SNAKE_CASE_   ,	eos_token=SCREAMING_SNAKE_CASE_   ,	unk_token=SCREAMING_SNAKE_CASE_   ,	sep_token=SCREAMING_SNAKE_CASE_   ,	cls_token=SCREAMING_SNAKE_CASE_   ,	pad_token=SCREAMING_SNAKE_CASE_   ,	mask_token=SCREAMING_SNAKE_CASE_   ,	sp_model_kwargs=self.sp_model_kwargs   ,	**SCREAMING_SNAKE_CASE_   ,	)

          lowerCamelCase__          =						spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(str(SCREAMING_SNAKE_CASE_							)							)
          lowerCamelCase__          =						vocab_file

          # Original fairseq vocab and spm vocab must be "aligned":
          # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
          # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
          # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
          # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

          # Mimic fairseq token-to-id alignment for the first 4 token
          lowerCamelCase__          =						{'<s>': 0, '<pad>': 1, '</s>': 2, '<unk>': 3}

          # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
          lowerCamelCase__          =						1

          lowerCamelCase__          =						len(self.sp_model							)
          lowerCamelCase__          =						{
              code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(SCREAMING_SNAKE_CASE_							)
          }
          lowerCamelCase__          =						{v: k for k, v in self.lang_code_to_id.items()}
          lowerCamelCase__          =						len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset

          self.fairseq_tokens_to_ids.update(self.lang_code_to_id							)
          lowerCamelCase__          =						{v: k for k, v in self.fairseq_tokens_to_ids.items()}

          lowerCamelCase__          =						src_lang if src_lang is not None else 'en_XX'
          lowerCamelCase__          =						self.lang_code_to_id[self._src_lang]
          lowerCamelCase__          =						tgt_lang
          self.set_src_lang_special_tokens(self._src_lang							)
   @property
   def    __UpperCAmelCase       (       self			:				Dict							):
          return len(self.sp_model							) + len(self.lang_code_to_id							) + self.fairseq_offset + 1  # Plus 1 for the mask token
   @property
   def    __UpperCAmelCase       (       self			:				int							):
          return self._src_lang
   @src_lang.setter
   def    __UpperCAmelCase       (       self			:				int   ,	SCREAMING_SNAKE_CASE_			:				str							):
          lowerCamelCase__          =						new_src_lang
          self.set_src_lang_special_tokens(self._src_lang							)
   def __getstate__(       self			:				Optional[Any]							):
          lowerCamelCase__          =						self.__dict__.copy()
          lowerCamelCase__          =						None
          return state
   def __setstate__(       self			:				Dict   ,	SCREAMING_SNAKE_CASE_			:				Dict							):
          lowerCamelCase__          =						d

          # for backward compatibility
          if not hasattr(self   ,	"""sp_model_kwargs"""							):
                 lowerCamelCase__          =						{}

          lowerCamelCase__          =						spm.SentencePieceProcessor(**self.sp_model_kwargs							)
          self.sp_model.Load(self.vocab_file							)
   def    __UpperCAmelCase       (       self			:				List[Any]							):
          lowerCamelCase__          =						{self.convert_ids_to_tokens(SCREAMING_SNAKE_CASE_							): i for i in range(self.vocab_size							)}
          vocab.update(self.added_tokens_encoder							)
          return vocab
   def    __UpperCAmelCase       (       self			:				List[Any]   ,	SCREAMING_SNAKE_CASE_			:				str							):
          return self.sp_model.encode(SCREAMING_SNAKE_CASE_   ,	out_type=SCREAMING_SNAKE_CASE_							)
   def    __UpperCAmelCase       (       self			:				Optional[int]   ,	SCREAMING_SNAKE_CASE_			:				str							):
          if token in self.fairseq_tokens_to_ids:
                 return self.fairseq_tokens_to_ids[token]
          lowerCamelCase__          =						self.sp_model.PieceToId(SCREAMING_SNAKE_CASE_							)

          # Need to return unknown token if the SP model returned 0
          return spm_id + self.fairseq_offset if spm_id else self.unk_token_id
   def    __UpperCAmelCase       (       self			:				List[Any]   ,	SCREAMING_SNAKE_CASE_			:				int							):
          if index in self.fairseq_ids_to_tokens:
                 return self.fairseq_ids_to_tokens[index]
          return self.sp_model.IdToPiece(index - self.fairseq_offset							)
   def    __UpperCAmelCase       (       self			:				Dict   ,	SCREAMING_SNAKE_CASE_			:				List[Any]							):
          lowerCamelCase__          =						[]
          lowerCamelCase__          =						''
          lowerCamelCase__          =						False
          for token in tokens:
                 # make sure that special tokens are not decoded using sentencepiece model
                 if token in self.all_special_tokens:
                        if not prev_is_special:
                               out_string += " "
                        out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_							) + token
                        lowerCamelCase__          =						True
                        lowerCamelCase__          =						[]
                 else:
                        current_sub_tokens.append(SCREAMING_SNAKE_CASE_							)
                        lowerCamelCase__          =						False
          out_string += self.sp_model.decode(SCREAMING_SNAKE_CASE_							)
          return out_string.strip()
   def    __UpperCAmelCase       (       self			:				Dict   ,	SCREAMING_SNAKE_CASE_			:				str   ,	SCREAMING_SNAKE_CASE_			:				Optional[str] = None							):
          if not os.path.isdir(SCREAMING_SNAKE_CASE_							):
                 logger.error(f"""Vocabulary path ({save_directory}) should be a directory"""							)
                 return
          lowerCamelCase__          =						os.path.join(
              SCREAMING_SNAKE_CASE_   ,	(filename_prefix + """-""" if filename_prefix else """""") + VOCAB_FILES_NAMES["""vocab_file"""]							)

          if os.path.abspath(self.vocab_file							) != os.path.abspath(SCREAMING_SNAKE_CASE_							) and os.path.isfile(self.vocab_file							):
                 copyfile(self.vocab_file   ,	SCREAMING_SNAKE_CASE_							)
          elif not os.path.isfile(self.vocab_file							):
                 with open(SCREAMING_SNAKE_CASE_   ,	"""wb"""							) as fi:
                        lowerCamelCase__          =						self.sp_model.serialized_model_proto()
                        fi.write(SCREAMING_SNAKE_CASE_							)

          return (out_vocab_file,)
   def    __UpperCAmelCase       (       self			:				Dict   ,	SCREAMING_SNAKE_CASE_			:				List[int]   ,	SCREAMING_SNAKE_CASE_			:				Optional[List[int]] = None   ,	SCREAMING_SNAKE_CASE_			:				bool = False							):

          if already_has_special_tokens:
                 return super().get_special_tokens_mask(
                     token_ids_a=SCREAMING_SNAKE_CASE_   ,	token_ids_a=SCREAMING_SNAKE_CASE_   ,	already_has_special_tokens=SCREAMING_SNAKE_CASE_							)

          lowerCamelCase__          =						[1] * len(self.prefix_tokens							)
          lowerCamelCase__          =						[1] * len(self.suffix_tokens							)
          if token_ids_a is None:
                 return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE_							)) + suffix_ones
          return prefix_ones + ([0] * len(SCREAMING_SNAKE_CASE_							)) + ([0] * len(SCREAMING_SNAKE_CASE_							)) + suffix_ones
   def    __UpperCAmelCase       (       self			:				Tuple   ,	SCREAMING_SNAKE_CASE_			:				List[int]   ,	SCREAMING_SNAKE_CASE_			:				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    __UpperCAmelCase       (       self			:				Union[str, Any]   ,	SCREAMING_SNAKE_CASE_			:				Optional[Any]   ,	SCREAMING_SNAKE_CASE_			:				str   ,	SCREAMING_SNAKE_CASE_			:				Optional[str]   ,	SCREAMING_SNAKE_CASE_			:				Optional[str]   ,	**SCREAMING_SNAKE_CASE_			:				List[str]							):
          if src_lang is None or tgt_lang is None:
                 raise ValueError("""Translation requires a `src_lang` and a `tgt_lang` for this model"""							)
          lowerCamelCase__          =						src_lang
          lowerCamelCase__          =						self(SCREAMING_SNAKE_CASE_   ,	add_special_tokens=SCREAMING_SNAKE_CASE_   ,	return_tensors=SCREAMING_SNAKE_CASE_   ,	**SCREAMING_SNAKE_CASE_							)
          lowerCamelCase__          =						self.convert_tokens_to_ids(SCREAMING_SNAKE_CASE_							)
          lowerCamelCase__          =						tgt_lang_id
          return inputs
   def    __UpperCAmelCase       (       self			:				Tuple   ,	SCREAMING_SNAKE_CASE_			:				List[str]   ,	SCREAMING_SNAKE_CASE_			:				str = "en_XX"   ,	SCREAMING_SNAKE_CASE_			:				Optional[List[str]] = None   ,	SCREAMING_SNAKE_CASE_			:				str = "ro_RO"   ,	**SCREAMING_SNAKE_CASE_			:				Optional[Any]   ,	):
          lowerCamelCase__          =						src_lang
          lowerCamelCase__          =						tgt_lang
          return super().prepare_seqaseq_batch(SCREAMING_SNAKE_CASE_   ,	SCREAMING_SNAKE_CASE_   ,	**SCREAMING_SNAKE_CASE_							)
   def    __UpperCAmelCase       (       self			:				Dict							):
          return self.set_src_lang_special_tokens(self.src_lang							)
   def    __UpperCAmelCase       (       self			:				List[Any]							):
          return self.set_tgt_lang_special_tokens(self.tgt_lang							)
   def    __UpperCAmelCase       (       self			:				int   ,	SCREAMING_SNAKE_CASE_			:				str							):
          lowerCamelCase__          =						self.lang_code_to_id[src_lang]
          lowerCamelCase__          =						[self.cur_lang_code_id]
          lowerCamelCase__          =						[self.eos_token_id]





   def    __UpperCAmelCase       (       self			:				List[str]   ,	SCREAMING_SNAKE_CASE_			:				str							):
          lowerCamelCase__          =						self.lang_code_to_id[tgt_lang]
          lowerCamelCase__          =						[self.cur_lang_code_id]
          lowerCamelCase__          =						[self.eos_token_id]














"""simple docstring"""


from typing import List, Union

from ..utils import (
    add_end_docstrings,
    is_tf_available,
    is_torch_available,
    is_vision_available,
    logging,
    requires_backends,
)
from .base import PIPELINE_INIT_ARGS, Pipeline


if is_vision_available():
		from PIL import Image

		from ..image_utils import load_image

if is_tf_available():
		from ..models.auto.modeling_tf_auto import TF_MODEL_FOR_VISION_2_SEQ_MAPPING

if is_torch_available():
		import torch

		from ..models.auto.modeling_auto import MODEL_FOR_VISION_2_SEQ_MAPPING

__SCREAMING_SNAKE_CASE   =   logging.get_logger(__name__)




@add_end_docstrings(UpperCamelCase						)
class 						__UpperCamelCase     (     UpperCamelCase						):
					def __init__(							self							:	List[Any]     , *UpperCAmelCase							:	Any     , **UpperCAmelCase							:	Tuple   )	->       Optional[int]:
											super().__init__(*UpperCAmelCase     , **UpperCAmelCase   )
											requires_backends(self     , 'vision'   )
											self.check_model_type(
											    TF_MODEL_FOR_VISION_2_SEQ_MAPPING if self.framework == 'tf' else MODEL_FOR_VISION_2_SEQ_MAPPING   )
					def        UpperCAmelCase__ (							self							:	List[str]     , UpperCAmelCase							:	str=None     , UpperCAmelCase							:	str=None     , UpperCAmelCase							:	List[str]=None   )	->       Tuple:
											lowerCAmelCase   :Optional[Any]						=			{}
											lowerCAmelCase   :int						=			{}

											if prompt is not None:
																	lowerCAmelCase   :Any						=			prompt

											if generate_kwargs is not None:
																	lowerCAmelCase   :List[Any]						=			generate_kwargs
											if max_new_tokens is not None:
																	if "generate_kwargs" not in forward_kwargs:
																							lowerCAmelCase   :Any						=			{}
																	if "max_new_tokens" in forward_kwargs["generate_kwargs"]:
																							raise ValueError(
																							    '\'max_new_tokens\' is defined twice, once in \'generate_kwargs\' and once as a direct parameter,'
																							    ' please use only one'   )
																	lowerCAmelCase   :Union[str, Any]						=			max_new_tokens
											return preprocess_params, forward_kwargs, {}
					def __call__(							self							:	Tuple     , UpperCAmelCase							:	Union[str, List[str], "Image.Image", List["Image.Image"]]     , **UpperCAmelCase							:	Any   )	->       Optional[int]:
											return super().__call__(UpperCAmelCase     , **UpperCAmelCase   )
					def        UpperCAmelCase__ (							self							:	List[str]     , UpperCAmelCase							:	Tuple     , UpperCAmelCase							:	Any=None   )	->       str:
											lowerCAmelCase   :Dict						=			load_image(UpperCAmelCase   )

											if prompt is not None:
																	if not isinstance(UpperCAmelCase     , UpperCAmelCase   ):
																							raise ValueError(
																							    f"""Received an invalid text input, got - {type(UpperCAmelCase   )} - but expected a single string. """
																							    'Note also that one single text can be provided for conditional image to text generation.'   )

																	lowerCAmelCase   :List[str]						=			self.model.config.model_type

																	if model_type == "git":
																							lowerCAmelCase   :Any						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )
																							lowerCAmelCase   :Union[str, Any]						=			self.tokenizer(text=UpperCAmelCase     , add_special_tokens=UpperCAmelCase   ).input_ids
																							lowerCAmelCase   :Optional[int]						=			[self.tokenizer.cls_token_id] + input_ids
																							lowerCAmelCase   :Union[str, Any]						=			torch.tensor(UpperCAmelCase   ).unsqueeze(0   )
																							model_inputs.update({'input_ids': input_ids}   )

																	elif model_type == "pix2struct":
																							lowerCAmelCase   :Optional[Any]						=			self.image_processor(images=UpperCAmelCase     , header_text=UpperCAmelCase     , return_tensors=self.framework   )

																	elif model_type != "vision-encoder-decoder":
																							# vision-encoder-decoder does not support conditional generation
																							lowerCAmelCase   :int						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )
																							lowerCAmelCase   :Optional[Any]						=			self.tokenizer(UpperCAmelCase     , return_tensors=self.framework   )
																							model_inputs.update(UpperCAmelCase   )

																	else:
																							raise ValueError(f"""Model type {model_type} does not support conditional text generation"""   )

											else:
																	lowerCAmelCase   :str						=			self.image_processor(images=UpperCAmelCase     , return_tensors=self.framework   )

											if self.model.config.model_type == "git" and prompt is None:
																	lowerCAmelCase   :Union[str, Any]						=			None

											return model_inputs
					def        UpperCAmelCase__ (							self							:	Optional[Any]     , UpperCAmelCase							:	Union[str, Any]     , UpperCAmelCase							:	Any=None   )	->       int:
											# Git model sets `model_inputs["input_ids"] = None` in `preprocess` (when `prompt=None`). In batch model, the
											# pipeline will group them into a list of `None`, which fail `_forward`. Avoid this by checking it first.
											if (
											    "input_ids" in model_inputs
											    and isinstance(model_inputs['input_ids']     , UpperCAmelCase   )
											    and all(x is None for x in model_inputs['input_ids']   )
											):
																	lowerCAmelCase   :Optional[int]						=			None

											if generate_kwargs is None:
																	lowerCAmelCase   :Optional[int]						=			{}
											# FIXME: We need to pop here due to a difference in how `generation.py` and `generation.tf_utils.py`
											#  parse inputs. In the Tensorflow version, `generate` raises an error if we don't use `input_ids` whereas
											#  the PyTorch version matches it with `self.model.main_input_name` or `self.model.encoder.main_input_name`
											#  in the `_prepare_model_inputs` method.
											lowerCAmelCase   :Optional[int]						=			model_inputs.pop(self.model.main_input_name   )
											lowerCAmelCase   :Any						=			self.model.generate(UpperCAmelCase     , **UpperCAmelCase     , **UpperCAmelCase   )
											return model_outputs
					def        UpperCAmelCase__ (							self							:	Dict     , UpperCAmelCase							:	Union[str, Any]   )	->       Any:
											lowerCAmelCase   :str						=			[]
											for output_ids in model_outputs:
																	lowerCAmelCase   :Optional[int]						=			{
																	    'generated_text': self.tokenizer.decode(
																	        UpperCAmelCase     , skip_special_tokens=UpperCAmelCase     , )
																	}
																	records.append(UpperCAmelCase   )
											return records