


'''simple docstring'''



# This model implementation is heavily inspired by https://github.com/haofanwang/ControlNet-for-Diffusers/

import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from PIL import Image
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import (
    AutoencoderKL,
    ControlNetModel,
    DDIMScheduler,
    StableDiffusionControlNetImgaImgPipeline,
    UNetaDConditionModel,
)
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import MultiControlNetModel
from diffusers.utils import floats_tensor, load_image, load_numpy, randn_tensor, slow, torch_device
from diffusers.utils.import_utils import is_xformers_available
from diffusers.utils.testing_utils import enable_full_determinism, require_torch_gpu

from ..pipeline_params import (
    IMAGE_TO_IMAGE_IMAGE_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS,
    TEXT_GUIDED_IMAGE_VARIATION_PARAMS,
)
from ..test_pipelines_common import (
    PipelineKarrasSchedulerTesterMixin,
    PipelineLatentTesterMixin,
    PipelineTesterMixin,
)


enable_full_determinism()






class  UpperCamelCase__(     lowerCAmelCase		,       lowerCAmelCase		,       lowerCAmelCase		,       unittest.TestCase ):
				__magic_name__       :   str							= StableDiffusionControlNetImgaImgPipeline
				__magic_name__       :   Optional[Any]							= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
				__magic_name__       :   int							= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				__magic_name__       :   str							= IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"} )
				__magic_name__       :   List[Any]							= IMAGE_TO_IMAGE_IMAGE_PARAMS




				def   a__(	self    :				Any   )->  Tuple:

										"""simple docstring"""




										torch.manual_seed(0   )
										UpperCAmelCase         =			UNetaDConditionModel(
										    block_out_channels=(32, 64)   ,					layers_per_block=2   ,					sample_size=32   ,					in_channels=4   ,					out_channels=4   ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,					up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,					cross_attention_dim=32   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			ControlNetModel(
										    block_out_channels=(32, 64)   ,					layers_per_block=2   ,					in_channels=4   ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,					cross_attention_dim=32   ,					conditioning_embedding_out_channels=(16, 32)   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			DDIMScheduler(
										    beta_start=0.00085   ,					beta_end=0.012   ,					beta_schedule='''scaled_linear'''   ,					clip_sample=lowerCAmelCase   ,					set_alpha_to_one=lowerCAmelCase   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			AutoencoderKL(
										    block_out_channels=[32, 64]   ,					in_channels=3   ,					out_channels=3   ,					down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,					up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,					latent_channels=4   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			CLIPTextConfig(
										    bos_token_id=0   ,					eos_token_id=2   ,					hidden_size=32   ,					intermediate_size=37   ,					layer_norm_eps=1E-05   ,					num_attention_heads=4   ,					num_hidden_layers=5   ,					pad_token_id=1   ,					vocab_size=1000   ,					)
										UpperCAmelCase         =			CLIPTextModel(lowerCAmelCase   )
										UpperCAmelCase         =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

										UpperCAmelCase         =			{
										    '''unet''': unet,
										    '''controlnet''': controlnet,
										    '''scheduler''': scheduler,
										    '''vae''': vae,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''safety_checker''': None,
										    '''feature_extractor''': None,
										}
										return components




				def   a__(	self    :				Any   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				List[str]=0   )->  Any:

										"""simple docstring"""




										if str(lowerCAmelCase   ).startswith('''mps'''   ):
																UpperCAmelCase         =			torch.manual_seed(lowerCAmelCase   )
										else:
																UpperCAmelCase         =			torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )

										UpperCAmelCase         =			2
										UpperCAmelCase         =			randn_tensor(
										    (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,					generator=lowerCAmelCase   ,					device=torch.device(lowerCAmelCase   )   ,					)
										UpperCAmelCase         =			floats_tensor(control_image.shape   ,					rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			image.cpu().permute(0   ,					2   ,					3   ,					1   )[0]
										UpperCAmelCase         =			Image.fromarray(np.uinta(lowerCAmelCase   )   ).convert('''RGB'''   ).resize((64, 64)   )
										UpperCAmelCase         =			{
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''generator''': generator,
										    '''num_inference_steps''': 2,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										    '''image''': image,
										    '''control_image''': control_image,
										}

										return inputs




				def   a__(	self    :				List[str]   )->  int:

										"""simple docstring"""




										return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3   )




				@unittest.skipIf(
				    torch_device != '''cuda''' or not is_xformers_available()   ,					reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,					)
				def   a__(	self    :				Dict   )->  str:

										"""simple docstring"""




										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3   )





				def   a__(	self    :				Optional[int]   )->  Optional[int]:

										"""simple docstring"""




										self._test_inference_batch_single_identical(expected_max_diff=2E-3   )







class  UpperCamelCase__(     lowerCAmelCase		,       lowerCAmelCase		,       unittest.TestCase ):
				__magic_name__       :   Tuple							= StableDiffusionControlNetImgaImgPipeline
				__magic_name__       :   str							= TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
				__magic_name__       :   str							= TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
				__magic_name__       :   Any							= frozenset([] )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess




				def   a__(	self    :				Dict   )->  List[Any]:

										"""simple docstring"""




										torch.manual_seed(0   )
										UpperCAmelCase         =			UNetaDConditionModel(
										    block_out_channels=(32, 64)   ,					layers_per_block=2   ,					sample_size=32   ,					in_channels=4   ,					out_channels=4   ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,					up_block_types=('''CrossAttnUpBlock2D''', '''UpBlock2D''')   ,					cross_attention_dim=32   ,					)
										torch.manual_seed(0   )

										def init_weights(lowerCAmelCase    :				Tuple   ):
																if isinstance(lowerCAmelCase   ,					torch.nn.Convad   ):
																						torch.nn.init.normal(m.weight   )
																						m.bias.data.fill_(1.0   )

										UpperCAmelCase         =			ControlNetModel(
										    block_out_channels=(32, 64)   ,					layers_per_block=2   ,					in_channels=4   ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,					cross_attention_dim=32   ,					conditioning_embedding_out_channels=(16, 32)   ,					)
										controlneta.controlnet_down_blocks.apply(lowerCAmelCase   )

										torch.manual_seed(0   )
										UpperCAmelCase         =			ControlNetModel(
										    block_out_channels=(32, 64)   ,					layers_per_block=2   ,					in_channels=4   ,					down_block_types=('''DownBlock2D''', '''CrossAttnDownBlock2D''')   ,					cross_attention_dim=32   ,					conditioning_embedding_out_channels=(16, 32)   ,					)
										controlneta.controlnet_down_blocks.apply(lowerCAmelCase   )

										torch.manual_seed(0   )
										UpperCAmelCase         =			DDIMScheduler(
										    beta_start=0.00085   ,					beta_end=0.012   ,					beta_schedule='''scaled_linear'''   ,					clip_sample=lowerCAmelCase   ,					set_alpha_to_one=lowerCAmelCase   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			AutoencoderKL(
										    block_out_channels=[32, 64]   ,					in_channels=3   ,					out_channels=3   ,					down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']   ,					up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']   ,					latent_channels=4   ,					)
										torch.manual_seed(0   )
										UpperCAmelCase         =			CLIPTextConfig(
										    bos_token_id=0   ,					eos_token_id=2   ,					hidden_size=32   ,					intermediate_size=37   ,					layer_norm_eps=1E-05   ,					num_attention_heads=4   ,					num_hidden_layers=5   ,					pad_token_id=1   ,					vocab_size=1000   ,					)
										UpperCAmelCase         =			CLIPTextModel(lowerCAmelCase   )
										UpperCAmelCase         =			CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''   )

										UpperCAmelCase         =			MultiControlNetModel([controlneta, controlneta]   )

										UpperCAmelCase         =			{
										    '''unet''': unet,
										    '''controlnet''': controlnet,
										    '''scheduler''': scheduler,
										    '''vae''': vae,
										    '''text_encoder''': text_encoder,
										    '''tokenizer''': tokenizer,
										    '''safety_checker''': None,
										    '''feature_extractor''': None,
										}
										return components




				def   a__(	self    :				List[str]   ,					lowerCAmelCase    :				int   ,					lowerCAmelCase    :				Union[str, Any]=0   )->  List[str]:

										"""simple docstring"""




										if str(lowerCAmelCase   ).startswith('''mps'''   ):
																UpperCAmelCase         =			torch.manual_seed(lowerCAmelCase   )
										else:
																UpperCAmelCase         =			torch.Generator(device=lowerCAmelCase   ).manual_seed(lowerCAmelCase   )

										UpperCAmelCase         =			2

										UpperCAmelCase         =			[
										    randn_tensor(
										        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,					generator=lowerCAmelCase   ,					device=torch.device(lowerCAmelCase   )   ,					),
										    randn_tensor(
										        (1, 3, 32 * controlnet_embedder_scale_factor, 32 * controlnet_embedder_scale_factor)   ,					generator=lowerCAmelCase   ,					device=torch.device(lowerCAmelCase   )   ,					),
										]

										UpperCAmelCase         =			floats_tensor(control_image[0].shape   ,					rng=random.Random(lowerCAmelCase   )   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			image.cpu().permute(0   ,					2   ,					3   ,					1   )[0]
										UpperCAmelCase         =			Image.fromarray(np.uinta(lowerCAmelCase   )   ).convert('''RGB'''   ).resize((64, 64)   )
										UpperCAmelCase         =			{
										    '''prompt''': '''A painting of a squirrel eating a burger''',
										    '''generator''': generator,
										    '''num_inference_steps''': 2,
										    '''guidance_scale''': 6.0,
										    '''output_type''': '''numpy''',
										    '''image''': image,
										    '''control_image''': control_image,
										}

										return inputs




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

										"""simple docstring"""




										UpperCAmelCase         =			self.get_dummy_components()
										UpperCAmelCase         =			self.pipeline_class(**lowerCAmelCase   )
										pipe.to(lowerCAmelCase   )

										UpperCAmelCase         =			10.0
										UpperCAmelCase         =			4

										UpperCAmelCase         =			self.get_dummy_inputs(lowerCAmelCase   )
										UpperCAmelCase         =			steps
										UpperCAmelCase         =			scale
										UpperCAmelCase         =			pipe(**lowerCAmelCase   )[0]

										UpperCAmelCase         =			self.get_dummy_inputs(lowerCAmelCase   )
										UpperCAmelCase         =			steps
										UpperCAmelCase         =			scale
										UpperCAmelCase         =			pipe(**lowerCAmelCase   ,					control_guidance_start=0.1   ,					control_guidance_end=0.2   )[0]

										UpperCAmelCase         =			self.get_dummy_inputs(lowerCAmelCase   )
										UpperCAmelCase         =			steps
										UpperCAmelCase         =			scale
										UpperCAmelCase         =			pipe(**lowerCAmelCase   ,					control_guidance_start=[0.1, 0.3]   ,					control_guidance_end=[0.2, 0.7]   )[0]

										UpperCAmelCase         =			self.get_dummy_inputs(lowerCAmelCase   )
										UpperCAmelCase         =			steps
										UpperCAmelCase         =			scale
										UpperCAmelCase         =			pipe(**lowerCAmelCase   ,					control_guidance_start=0.4   ,					control_guidance_end=[0.5, 0.8]   )[0]

										# make sure that all outputs are different
										assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3
										assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3
										assert np.sum(np.abs(output_a - output_a   )   ) > 1E-3




				def   a__(	self    :				str   )->  List[str]:

										"""simple docstring"""




										return self._test_attention_slicing_forward_pass(expected_max_diff=2E-3   )




				@unittest.skipIf(
				    torch_device != '''cuda''' or not is_xformers_available()   ,					reason='''XFormers attention is only available with CUDA and `xformers` installed'''   ,					)
				def   a__(	self    :				Tuple   )->  Tuple:

										"""simple docstring"""




										self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3   )




				def   a__(	self    :				int   )->  Optional[int]:

										"""simple docstring"""




										self._test_inference_batch_single_identical(expected_max_diff=2E-3   )





				def   a__(	self    :				int   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			self.get_dummy_components()
										UpperCAmelCase         =			self.pipeline_class(**lowerCAmelCase   )
										pipe.to(lowerCAmelCase   )
										pipe.set_progress_bar_config(disable=lowerCAmelCase   )
										with tempfile.TemporaryDirectory() as tmpdir:
																try:
																						# save_pretrained is not implemented for Multi-ControlNet
																						pipe.save_pretrained(lowerCAmelCase   )
																except NotImplementedError:
																						pass







@slow
@require_torch_gpu
class  UpperCamelCase__(     unittest.TestCase ):




				def   a__(	self    :				Optional[Any]   )->  int:

										"""simple docstring"""




										super().tearDown()
										gc.collect()
										torch.cuda.empty_cache()





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

										"""simple docstring"""




										UpperCAmelCase         =			ControlNetModel.from_pretrained('''lllyasviel/sd-controlnet-canny'''   )

										UpperCAmelCase         =			StableDiffusionControlNetImgaImgPipeline.from_pretrained(
										    '''runwayml/stable-diffusion-v1-5'''   ,					safety_checker=lowerCAmelCase   ,					controlnet=lowerCAmelCase   )
										pipe.enable_model_cpu_offload()
										pipe.set_progress_bar_config(disable=lowerCAmelCase   )

										UpperCAmelCase         =			torch.Generator(device='''cpu'''   ).manual_seed(0   )
										UpperCAmelCase         =			'''evil space-punk bird'''
										UpperCAmelCase         =			load_image(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/bird_canny.png'''   ).resize((512, 512)   )
										UpperCAmelCase         =			load_image(
										    '''https://huggingface.co/lllyasviel/sd-controlnet-canny/resolve/main/images/bird.png'''   ).resize((512, 512)   )

										UpperCAmelCase         =			pipe(
										    lowerCAmelCase   ,					lowerCAmelCase   ,					control_image=lowerCAmelCase   ,					generator=lowerCAmelCase   ,					output_type='''np'''   ,					num_inference_steps=50   ,					strength=0.6   ,					)

										UpperCAmelCase         =			output.images[0]

										assert image.shape == (512, 512, 3)

										UpperCAmelCase         =			load_numpy(
										    '''https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd_controlnet/img2img.npy'''   )

										assert np.abs(expected_image - image   ).max() < 9E-2









'''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


_lowercase    :    str			     =  logging.get_logger(__name__)

_lowercase    :    List[Any]			     =  """▁"""

_lowercase    :    Optional[Any]			     =  {"""vocab_file""": """sentencepiece.bpe.model"""}

_lowercase    :    Any			     =  {
    """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"""
        ),
    }
}

_lowercase    :    int			     =  {
    """facebook/mbart-large-50-one-to-many-mmt""": 1024,
}

# fmt: off
_lowercase    :    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__(     lowerCAmelCase ):
				__magic_name__       :   Optional[int]							= VOCAB_FILES_NAMES
				__magic_name__       :   Union[str, Any]							= PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
				__magic_name__       :   Optional[int]							= PRETRAINED_VOCAB_FILES_MAP
				__magic_name__       :   Union[str, Any]							= ["input_ids", "attention_mask"]

				__magic_name__       :   List[int]							= []
				__magic_name__       :   List[int]							= []




				def __init__(	self    :				List[str]   ,					lowerCAmelCase    :				Optional[int]   ,					lowerCAmelCase    :				Optional[int]=None   ,					lowerCAmelCase    :				str=None   ,					lowerCAmelCase    :				List[Any]="</s>"   ,					lowerCAmelCase    :				Union[str, Any]="</s>"   ,					lowerCAmelCase    :				Union[str, Any]="<s>"   ,					lowerCAmelCase    :				int="<unk>"   ,					lowerCAmelCase    :				str="<pad>"   ,					lowerCAmelCase    :				Optional[int]="<mask>"   ,					lowerCAmelCase    :				Optional[Dict[str, Any]] = None   ,					**lowerCAmelCase    :				List[Any]   ,					)->  None:

										"""simple docstring"""




										UpperCAmelCase         =			AddedToken(lowerCAmelCase   ,					lstrip=lowerCAmelCase   ,					rstrip=lowerCAmelCase   ) if isinstance(lowerCAmelCase   ,					lowerCAmelCase   ) else mask_token

										UpperCAmelCase         =			{} if sp_model_kwargs is None else sp_model_kwargs

										UpperCAmelCase         =			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=lowerCAmelCase   ,					tgt_lang=lowerCAmelCase   ,					eos_token=lowerCAmelCase   ,					unk_token=lowerCAmelCase   ,					sep_token=lowerCAmelCase   ,					cls_token=lowerCAmelCase   ,					pad_token=lowerCAmelCase   ,					mask_token=lowerCAmelCase   ,					sp_model_kwargs=self.sp_model_kwargs   ,					**lowerCAmelCase   ,					)

										UpperCAmelCase         =			spm.SentencePieceProcessor(**self.sp_model_kwargs   )
										self.sp_model.Load(str(lowerCAmelCase   )   )
										UpperCAmelCase         =			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
										UpperCAmelCase         =			{'''<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
										UpperCAmelCase         =			1

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

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

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




				@property
				def   a__(	self    :				Union[str, Any]   )->  int:

										"""simple docstring"""




										return len(self.sp_model   ) + len(self.lang_code_to_id   ) + self.fairseq_offset + 1  # Plus 1 for the mask token




				@property
				def   a__(	self    :				str   )->  str:

										"""simple docstring"""




										return self._src_lang




				@src_lang.setter
				def   a__(	self    :				Any   ,					lowerCAmelCase    :				str   )->  None:

										"""simple docstring"""




										UpperCAmelCase         =			new_src_lang
										self.set_src_lang_special_tokens(self._src_lang   )




				def __getstate__(	self    :				Tuple   )->  Dict:

										"""simple docstring"""




										UpperCAmelCase         =			self.__dict__.copy()
										UpperCAmelCase         =			None
										return state




				def __setstate__(	self    :				Dict   ,					lowerCAmelCase    :				Dict   )->  None:

										"""simple docstring"""




										UpperCAmelCase         =			d

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

										UpperCAmelCase         =			spm.SentencePieceProcessor(**self.sp_model_kwargs   )
										self.sp_model.Load(self.vocab_file   )




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

										"""simple docstring"""




										UpperCAmelCase         =			{self.convert_ids_to_tokens(lowerCAmelCase   ): i for i in range(self.vocab_size   )}
										vocab.update(self.added_tokens_encoder   )
										return vocab




				def   a__(	self    :				str   ,					lowerCAmelCase    :				str   )->  List[str]:

										"""simple docstring"""




										return self.sp_model.encode(lowerCAmelCase   ,					out_type=lowerCAmelCase   )




				def   a__(	self    :				Optional[int]   ,					lowerCAmelCase    :				str   )->  int:

										"""simple docstring"""




										if token in self.fairseq_tokens_to_ids:
																return self.fairseq_tokens_to_ids[token]
										UpperCAmelCase         =			self.sp_model.PieceToId(lowerCAmelCase   )

										# 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   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				int   )->  str:

										"""simple docstring"""




										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   a__(	self    :				int   ,					lowerCAmelCase    :				List[Any]   )->  Dict:

										"""simple docstring"""




										UpperCAmelCase         =			[]
										UpperCAmelCase         =			''''''
										UpperCAmelCase         =			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(lowerCAmelCase   ) + token
																						UpperCAmelCase         =			True
																						UpperCAmelCase         =			[]
																else:
																						current_sub_tokens.append(lowerCAmelCase   )
																						UpperCAmelCase         =			False
										out_string += self.sp_model.decode(lowerCAmelCase   )
										return out_string.strip()




				def   a__(	self    :				Union[str, Any]   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				Optional[str] = None   )->  Tuple[str]:

										"""simple docstring"""




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

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

										return (out_vocab_file,)




				def   a__(	self    :				List[str]   ,					lowerCAmelCase    :				List[int]   ,					lowerCAmelCase    :				Optional[List[int]] = None   ,					lowerCAmelCase    :				bool = False   )->  List[int]:

										"""simple docstring"""




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

										UpperCAmelCase         =			[1] * len(self.prefix_tokens   )
										UpperCAmelCase         =			[1] * len(self.suffix_tokens   )
										if token_ids_a is None:
																return prefix_ones + ([0] * len(lowerCAmelCase   )) + suffix_ones
										return prefix_ones + ([0] * len(lowerCAmelCase   )) + ([0] * len(lowerCAmelCase   )) + suffix_ones




				def   a__(	self    :				Tuple   ,					lowerCAmelCase    :				List[int]   ,					lowerCAmelCase    :				Optional[List[int]] = 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   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				Union[str, Any]   ,					lowerCAmelCase    :				str   ,					lowerCAmelCase    :				Optional[str]   ,					lowerCAmelCase    :				Optional[str]   ,					**lowerCAmelCase    :				Optional[int]   )->  Optional[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'''   )
										UpperCAmelCase         =			src_lang
										UpperCAmelCase         =			self(lowerCAmelCase   ,					add_special_tokens=lowerCAmelCase   ,					return_tensors=lowerCAmelCase   ,					**lowerCAmelCase   )
										UpperCAmelCase         =			self.convert_tokens_to_ids(lowerCAmelCase   )
										UpperCAmelCase         =			tgt_lang_id
										return inputs




				def   a__(	self    :				str   ,					lowerCAmelCase    :				List[str]   ,					lowerCAmelCase    :				str = "en_XX"   ,					lowerCAmelCase    :				Optional[List[str]] = None   ,					lowerCAmelCase    :				str = "ro_RO"   ,					**lowerCAmelCase    :				List[str]   ,					)->  BatchEncoding:

										"""simple docstring"""




										UpperCAmelCase         =			src_lang
										UpperCAmelCase         =			tgt_lang
										return super().prepare_seqaseq_batch(lowerCAmelCase   ,					lowerCAmelCase   ,					**lowerCAmelCase   )




				def   a__(	self    :				Optional[int]   )->  Union[str, Any]:

										"""simple docstring"""




										return self.set_src_lang_special_tokens(self.src_lang   )




				def   a__(	self    :				List[Any]   )->  int:

										"""simple docstring"""




										return self.set_tgt_lang_special_tokens(self.tgt_lang   )




				def   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				str   )->  None:

										"""simple docstring"""




										UpperCAmelCase         =			self.lang_code_to_id[src_lang]
										UpperCAmelCase         =			[self.cur_lang_code_id]
										UpperCAmelCase         =			[self.eos_token_id]





				def   a__(	self    :				int   ,					lowerCAmelCase    :				str   )->  None:

										"""simple docstring"""




										UpperCAmelCase         =			self.lang_code_to_id[tgt_lang]
										UpperCAmelCase         =			[self.cur_lang_code_id]
										UpperCAmelCase         =			[self.eos_token_id]





