





# 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 						lowerCamelCase       (__lowerCamelCase							,     __lowerCamelCase							,     __lowerCamelCase							,     unittest.TestCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   StableDiffusionControlNetImgaImgPipeline
							UpperCAmelCase_      						=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
							UpperCAmelCase_      						=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							UpperCAmelCase_      						=   IMAGE_TO_IMAGE_IMAGE_PARAMS.union({"control_image"}      )
							UpperCAmelCase_      						=   IMAGE_TO_IMAGE_IMAGE_PARAMS






							def  A_     ( self     :					Tuple						)				->			int:





														"""simple docstring"""
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      str     =						UNetaDConditionModel(
														    block_out_channels=(3_2, 6_4),			layers_per_block=2,			sample_size=3_2,			in_channels=4,			out_channels=4,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),			cross_attention_dim=3_2,			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						ControlNetModel(
														    block_out_channels=(3_2, 6_4),			layers_per_block=2,			in_channels=4,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),			cross_attention_dim=3_2,			conditioning_embedding_out_channels=(1_6, 3_2),			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						DDIMScheduler(
														    beta_start=0.00085,			beta_end=0.012,			beta_schedule="scaled_linear",			clip_sample=_UpperCAmelCase,			set_alpha_to_one=_UpperCAmelCase,			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Any     =						AutoencoderKL(
														    block_out_channels=[3_2, 6_4],			in_channels=3,			out_channels=3,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],			latent_channels=4,			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						CLIPTextConfig(
														    bos_token_id=0,			eos_token_id=2,			hidden_size=3_2,			intermediate_size=3_7,			layer_norm_eps=1E-05,			num_attention_heads=4,			num_hidden_layers=5,			pad_token_id=1,			vocab_size=1_0_0_0,			)
														SCREAMING_SNAKE_CASE__       :      int     =						CLIPTextModel(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      int     =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

														SCREAMING_SNAKE_CASE__       :      int     =						{
														    "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     :					int,			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					str=0						)				->			int:





														"""simple docstring"""
														if str(_UpperCAmelCase						).startswith("mps"						):
																					SCREAMING_SNAKE_CASE__       :      Optional[int]     =						torch.manual_seed(_UpperCAmelCase						)
														else:
																					SCREAMING_SNAKE_CASE__       :      int     =						torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      str     =						2
														SCREAMING_SNAKE_CASE__       :      str     =						randn_tensor(
														    (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),			generator=_UpperCAmelCase,			device=torch.device(_UpperCAmelCase						),			)
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						floats_tensor(control_image.shape,			rng=random.Random(_UpperCAmelCase						)						).to(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Any     =						image.cpu().permute(0,			2,			3,			1						)[0]
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						Image.fromarray(np.uinta(_UpperCAmelCase						)						).convert("RGB"						).resize((6_4, 6_4)						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						{
														    "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     :					Any						)				->			Union[str, Any]:





														"""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     :					List[str]						)				->			Dict:





														"""simple docstring"""
														self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3						)






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





														"""simple docstring"""
														self._test_inference_batch_single_identical(expected_max_diff=2E-3						)


class 						lowerCamelCase       (__lowerCamelCase							,     __lowerCamelCase							,     unittest.TestCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   StableDiffusionControlNetImgaImgPipeline
							UpperCAmelCase_      						=   TEXT_GUIDED_IMAGE_VARIATION_PARAMS - {"height", "width"}
							UpperCAmelCase_      						=   TEXT_GUIDED_IMAGE_VARIATION_BATCH_PARAMS
							UpperCAmelCase_      						=   frozenset([]      )  # TO_DO: add image_params once refactored VaeImageProcessor.preprocess






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





														"""simple docstring"""
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      str     =						UNetaDConditionModel(
														    block_out_channels=(3_2, 6_4),			layers_per_block=2,			sample_size=3_2,			in_channels=4,			out_channels=4,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),			up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),			cross_attention_dim=3_2,			)
														torch.manual_seed(0						)

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

														SCREAMING_SNAKE_CASE__       :      Tuple     =						ControlNetModel(
														    block_out_channels=(3_2, 6_4),			layers_per_block=2,			in_channels=4,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),			cross_attention_dim=3_2,			conditioning_embedding_out_channels=(1_6, 3_2),			)
														controlneta.controlnet_down_blocks.apply(_UpperCAmelCase						)

														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						ControlNetModel(
														    block_out_channels=(3_2, 6_4),			layers_per_block=2,			in_channels=4,			down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),			cross_attention_dim=3_2,			conditioning_embedding_out_channels=(1_6, 3_2),			)
														controlneta.controlnet_down_blocks.apply(_UpperCAmelCase						)

														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						DDIMScheduler(
														    beta_start=0.00085,			beta_end=0.012,			beta_schedule="scaled_linear",			clip_sample=_UpperCAmelCase,			set_alpha_to_one=_UpperCAmelCase,			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Any     =						AutoencoderKL(
														    block_out_channels=[3_2, 6_4],			in_channels=3,			out_channels=3,			down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],			up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],			latent_channels=4,			)
														torch.manual_seed(0						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						CLIPTextConfig(
														    bos_token_id=0,			eos_token_id=2,			hidden_size=3_2,			intermediate_size=3_7,			layer_norm_eps=1E-05,			num_attention_heads=4,			num_hidden_layers=5,			pad_token_id=1,			vocab_size=1_0_0_0,			)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						CLIPTextModel(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"						)

														SCREAMING_SNAKE_CASE__       :      List[str]     =						MultiControlNetModel([controlneta, controlneta]						)

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						{
														    "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     :					Optional[int],			_UpperCAmelCase     :					Any,			_UpperCAmelCase     :					List[str]=0						)				->			List[Any]:





														"""simple docstring"""
														if str(_UpperCAmelCase						).startswith("mps"						):
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						torch.manual_seed(_UpperCAmelCase						)
														else:
																					SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						torch.Generator(device=_UpperCAmelCase						).manual_seed(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						2

														SCREAMING_SNAKE_CASE__       :      List[str]     =						[
														    randn_tensor(
														        (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),			generator=_UpperCAmelCase,			device=torch.device(_UpperCAmelCase						),			),
														    randn_tensor(
														        (1, 3, 3_2 * controlnet_embedder_scale_factor, 3_2 * controlnet_embedder_scale_factor),			generator=_UpperCAmelCase,			device=torch.device(_UpperCAmelCase						),			),
														]

														SCREAMING_SNAKE_CASE__       :      str     =						floats_tensor(control_image[0].shape,			rng=random.Random(_UpperCAmelCase						)						).to(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      str     =						image.cpu().permute(0,			2,			3,			1						)[0]
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						Image.fromarray(np.uinta(_UpperCAmelCase						)						).convert("RGB"						).resize((6_4, 6_4)						)
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						{
														    "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     :					Tuple						)				->			List[str]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						self.get_dummy_components()
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						self.pipeline_class(**_UpperCAmelCase						)
														pipe.to(_UpperCAmelCase						)

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						10.0
														SCREAMING_SNAKE_CASE__       :      Tuple     =						4

														SCREAMING_SNAKE_CASE__       :      str     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      List[str]     =						steps
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						scale
														SCREAMING_SNAKE_CASE__       :      int     =						pipe(**_UpperCAmelCase						)[0]

														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Dict     =						steps
														SCREAMING_SNAKE_CASE__       :      Dict     =						scale
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						pipe(**_UpperCAmelCase,			control_guidance_start=0.1,			control_guidance_end=0.2						)[0]

														SCREAMING_SNAKE_CASE__       :      str     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      int     =						steps
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						scale
														SCREAMING_SNAKE_CASE__       :      Tuple     =						pipe(**_UpperCAmelCase,			control_guidance_start=[0.1, 0.3],			control_guidance_end=[0.2, 0.7]						)[0]

														SCREAMING_SNAKE_CASE__       :      Tuple     =						self.get_dummy_inputs(_UpperCAmelCase						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						steps
														SCREAMING_SNAKE_CASE__       :      List[str]     =						scale
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						pipe(**_UpperCAmelCase,			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     :					List[str]						)				->			Dict:





														"""simple docstring"""
														self._test_xformers_attention_forwardGenerator_pass(expected_max_diff=2E-3						)






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





														"""simple docstring"""
														self._test_inference_batch_single_identical(expected_max_diff=2E-3						)






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.get_dummy_components()
														SCREAMING_SNAKE_CASE__       :      Any     =						self.pipeline_class(**_UpperCAmelCase						)
														pipe.to(_UpperCAmelCase						)
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)
														with tempfile.TemporaryDirectory() as tmpdir:
																					try:
																												# save_pretrained is not implemented for Multi-ControlNet
																												pipe.save_pretrained(_UpperCAmelCase						)
																					except NotImplementedError:
																												pass


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







							"""simple docstring"""




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





														"""simple docstring"""
														super().tearDown()
														gc.collect()
														torch.cuda.empty_cache()






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Any     =						ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny"						)

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						StableDiffusionControlNetImgaImgPipeline.from_pretrained(
														    "runwayml/stable-diffusion-v1-5",			safety_checker=_UpperCAmelCase,			controlnet=_UpperCAmelCase						)
														pipe.enable_model_cpu_offload()
														pipe.set_progress_bar_config(disable=_UpperCAmelCase						)

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

														SCREAMING_SNAKE_CASE__       :      List[Any]     =						pipe(
														    _UpperCAmelCase,			_UpperCAmelCase,			control_image=_UpperCAmelCase,			generator=_UpperCAmelCase,			output_type="np",			num_inference_steps=5_0,			strength=0.6,			)

														SCREAMING_SNAKE_CASE__       :      Dict     =						output.images[0]

														assert image.shape == (5_1_2, 5_1_2, 3)

														SCREAMING_SNAKE_CASE__       :      Tuple     =						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













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

import sentencepiece as spm

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


_lowerCamelCase :    int  							= logging.get_logger(__name__)

_lowerCamelCase :    Optional[Any]  							= '''▁'''

_lowerCamelCase :    Dict  							= {'''vocab_file''': '''sentencepiece.bpe.model'''}

_lowerCamelCase :    int  							= {
    '''vocab_file''': {
        '''xlm-roberta-base''': '''https://huggingface.co/xlm-roberta-base/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large''': '''https://huggingface.co/xlm-roberta-large/resolve/main/sentencepiece.bpe.model''',
        '''xlm-roberta-large-finetuned-conll02-dutch''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-dutch/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll02-spanish''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll02-spanish/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-english''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-english/resolve/main/sentencepiece.bpe.model'''
        ),
        '''xlm-roberta-large-finetuned-conll03-german''': (
            '''https://huggingface.co/xlm-roberta-large-finetuned-conll03-german/resolve/main/sentencepiece.bpe.model'''
        ),
    }
}

_lowerCamelCase :    Optional[Any]  							= {
    '''xlm-roberta-base''': 5_1_2,
    '''xlm-roberta-large''': 5_1_2,
    '''xlm-roberta-large-finetuned-conll02-dutch''': 5_1_2,
    '''xlm-roberta-large-finetuned-conll02-spanish''': 5_1_2,
    '''xlm-roberta-large-finetuned-conll03-english''': 5_1_2,
    '''xlm-roberta-large-finetuned-conll03-german''': 5_1_2,
}


class 						lowerCamelCase       (__lowerCamelCase      ):







							"""simple docstring"""




							UpperCAmelCase_      						=   VOCAB_FILES_NAMES
							UpperCAmelCase_      						=   PRETRAINED_VOCAB_FILES_MAP
							UpperCAmelCase_      						=   PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							UpperCAmelCase_      						=   ["input_ids", "attention_mask"]






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





														"""simple docstring"""
														# Mask token behave like a normal word, i.e. include the space before it
														SCREAMING_SNAKE_CASE__       :      int     =						AddedToken(_UpperCAmelCase,			lstrip=_UpperCAmelCase,			rstrip=_UpperCAmelCase						) if isinstance(_UpperCAmelCase,			_UpperCAmelCase						) else mask_token

														SCREAMING_SNAKE_CASE__       :      List[str]     =						{} if sp_model_kwargs is None else sp_model_kwargs

														super().__init__(
														    bos_token=_UpperCAmelCase,			eos_token=_UpperCAmelCase,			unk_token=_UpperCAmelCase,			sep_token=_UpperCAmelCase,			cls_token=_UpperCAmelCase,			pad_token=_UpperCAmelCase,			mask_token=_UpperCAmelCase,			sp_model_kwargs=self.sp_model_kwargs,			**_UpperCAmelCase,			)

														SCREAMING_SNAKE_CASE__       :      str     =						spm.SentencePieceProcessor(**self.sp_model_kwargs						)
														self.sp_model.Load(str(_UpperCAmelCase						)						)
														SCREAMING_SNAKE_CASE__       :      Tuple     =						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
														SCREAMING_SNAKE_CASE__       :      List[str]     =						{"<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
														SCREAMING_SNAKE_CASE__       :      Dict     =						1

														SCREAMING_SNAKE_CASE__       :      int     =						len(self.sp_model						) + self.fairseq_offset
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						{v: k for k, v in self.fairseq_tokens_to_ids.items()}






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						self.__dict__.copy()
														SCREAMING_SNAKE_CASE__       :      List[Any]     =						None
														SCREAMING_SNAKE_CASE__       :      Dict     =						self.sp_model.serialized_model_proto()
														return state






							def __setstate__( self     :					int,			_UpperCAmelCase     :					List[Any]						)				->			Any:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      Tuple     =						d

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

														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						spm.SentencePieceProcessor(**self.sp_model_kwargs						)
														self.sp_model.LoadFromSerializedProto(self.sp_model_proto						)






							def  A_     ( self     :					Any,			_UpperCAmelCase     :					List[int],			_UpperCAmelCase     :					Optional[List[int]] = None						)				->			List[int]:





														"""simple docstring"""

														if token_ids_a is None:
																					return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
														SCREAMING_SNAKE_CASE__       :      List[str]     =						[self.cls_token_id]
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						[self.sep_token_id]
														return cls + token_ids_a + sep + sep + token_ids_a + sep






							def  A_     ( self     :					List[Any],			_UpperCAmelCase     :					List[int],			_UpperCAmelCase     :					Optional[List[int]] = None,			_UpperCAmelCase     :					bool = False						)				->			List[int]:





														"""simple docstring"""

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

														if token_ids_a is None:
																					return [1] + ([0] * len(_UpperCAmelCase						)) + [1]
														return [1] + ([0] * len(_UpperCAmelCase						)) + [1, 1] + ([0] * len(_UpperCAmelCase						)) + [1]






							def  A_     ( self     :					Union[str, Any],			_UpperCAmelCase     :					List[int],			_UpperCAmelCase     :					Optional[List[int]] = None						)				->			List[int]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						[self.sep_token_id]
														SCREAMING_SNAKE_CASE__       :      Optional[int]     =						[self.cls_token_id]

														if token_ids_a is None:
																					return len(cls + token_ids_a + sep						) * [0]
														return len(cls + token_ids_a + sep + sep + token_ids_a + sep						) * [0]






							@property
							def  A_     ( self     :					List[str]						)				->			List[str]:





														"""simple docstring"""
														return len(self.sp_model						) + self.fairseq_offset + 1  # Add the <mask> token






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





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      str     =						{self.convert_ids_to_tokens(_UpperCAmelCase						): i for i in range(self.vocab_size						)}
														vocab.update(self.added_tokens_encoder						)
														return vocab






							def  A_     ( self     :					List[str],			_UpperCAmelCase     :					str						)				->			List[str]:





														"""simple docstring"""
														return self.sp_model.encode(_UpperCAmelCase,			out_type=_UpperCAmelCase						)






							def  A_     ( self     :					Optional[Any],			_UpperCAmelCase     :					Union[str, Any]						)				->			Union[str, Any]:





														"""simple docstring"""
														if token in self.fairseq_tokens_to_ids:
																					return self.fairseq_tokens_to_ids[token]
														SCREAMING_SNAKE_CASE__       :      Union[str, Any]     =						self.sp_model.PieceToId(_UpperCAmelCase						)

														# 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     :					Tuple,			_UpperCAmelCase     :					List[str]						)				->			List[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     :					Any,			_UpperCAmelCase     :					int						)				->			List[Any]:





														"""simple docstring"""
														SCREAMING_SNAKE_CASE__       :      int     =						"".join(_UpperCAmelCase						).replace(_UpperCAmelCase,			" "						).strip()
														return out_string






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





														"""simple docstring"""
														if not os.path.isdir(_UpperCAmelCase						):
																					logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''						)
																					return
														SCREAMING_SNAKE_CASE__       :      Optional[Any]     =						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						) and os.path.isfile(self.vocab_file						):
																					copyfile(self.vocab_file,			_UpperCAmelCase						)
														elif not os.path.isfile(self.vocab_file						):
																					with open(_UpperCAmelCase,			"wb"						) as fi:
																												SCREAMING_SNAKE_CASE__       :      Any     =						self.sp_model.serialized_model_proto()
																												fi.write(_UpperCAmelCase						)

														return (out_vocab_file,)






