

"""simple docstring"""




from typing import List, Optional, Union

import numpy as np
import PIL
import torch
from PIL import Image

from ...models import UNetaDConditionModel, VQModel
from ...pipelines import DiffusionPipeline
from ...pipelines.pipeline_utils import ImagePipelineOutput
from ...schedulers import DDPMScheduler
from ...utils import (
    is_accelerate_available,
    is_accelerate_version,
    logging,
    randn_tensor,
    replace_example_docstring,
)


_UpperCamelCase			      =		logging.get_logger(__name__)  # pylint: disable=invalid-name

_UpperCamelCase			      =		"\n    Examples:\n        ```py\n        >>> import torch\n        >>> import numpy as np\n\n        >>> from diffusers import KandinskyV22PriorEmb2EmbPipeline, KandinskyV22ControlnetImg2ImgPipeline\n        >>> from transformers import pipeline\n        >>> from diffusers.utils import load_image\n\n\n        >>> def make_hint(image, depth_estimator):\n        ...     image = depth_estimator(image)[\"depth\"]\n        ...     image = np.array(image)\n        ...     image = image[:, :, None]\n        ...     image = np.concatenate([image, image, image], axis=2)\n        ...     detected_map = torch.from_numpy(image).float() / 255.0\n        ...     hint = detected_map.permute(2, 0, 1)\n        ...     return hint\n\n\n        >>> depth_estimator = pipeline(\"depth-estimation\")\n\n        >>> pipe_prior = KandinskyV22PriorEmb2EmbPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-prior\", torch_dtype=torch.float16\n        ... )\n        >>> pipe_prior = pipe_prior.to(\"cuda\")\n\n        >>> pipe = KandinskyV22ControlnetImg2ImgPipeline.from_pretrained(\n        ...     \"kandinsky-community/kandinsky-2-2-controlnet-depth\", torch_dtype=torch.float16\n        ... )\n        >>> pipe = pipe.to(\"cuda\")\n\n        >>> img = load_image(\n        ...     \"https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main\"\n        ...     \"/kandinsky/cat.png\"\n        ... ).resize((768, 768))\n\n\n        >>> hint = make_hint(img, depth_estimator).unsqueeze(0).half().to(\"cuda\")\n\n        >>> prompt = \"A robot, 4k photo\"\n        >>> negative_prior_prompt = \"lowres, text, error, cropped, worst quality, low quality, jpeg artifacts, ugly, duplicate, morbid, mutilated, out of frame, extra fingers, mutated hands, poorly drawn hands, poorly drawn face, mutation, deformed, blurry, dehydrated, bad anatomy, bad proportions, extra limbs, cloned face, disfigured, gross proportions, malformed limbs, missing arms, missing legs, extra arms, extra legs, fused fingers, too many fingers, long neck, username, watermark, signature\"\n\n        >>> generator = torch.Generator(device=\"cuda\").manual_seed(43)\n\n        >>> img_emb = pipe_prior(prompt=prompt, image=img, strength=0.85, generator=generator)\n        >>> negative_emb = pipe_prior(prompt=negative_prior_prompt, image=img, strength=1, generator=generator)\n\n        >>> images = pipe(\n        ...     image=img,\n        ...     strength=0.5,\n        ...     image_embeds=img_emb.image_embeds,\n        ...     negative_image_embeds=negative_emb.image_embeds,\n        ...     hint=hint,\n        ...     num_inference_steps=50,\n        ...     generator=generator,\n        ...     height=768,\n        ...     width=768,\n        ... ).images\n\n        >>> images[0].save(\"robot_cat.png\")\n        ```\n"







def 		_lowercase			(					lowercase__       ,     lowercase__       ,     lowercase__=8						):
      __lowerCAmelCase :						Optional[int] = height // scale_factor**2
      if height % scale_factor**2 != 0:
            new_height += 1
      __lowerCAmelCase :						Dict = width // scale_factor**2
      if width % scale_factor**2 != 0:
            new_width += 1
      return new_height * scale_factor, new_width * scale_factor







def 		_lowercase			(					lowercase__       ,     lowercase__=5_1_2       ,     lowercase__=5_1_2						):
      __lowerCAmelCase :						Tuple = pil_image.resize((w, h)       ,     resample=Image.BICUBIC       ,     reducing_gap=1						)
      __lowerCAmelCase :						Optional[Any] = np.array(pil_image.convert('''RGB'''						)						)
      __lowerCAmelCase :						List[str] = arr.astype(np.floataa						) / 1_2_7.5 - 1
      __lowerCAmelCase :						Union[str, Any] = np.transpose(lowercase__       ,     [2, 0, 1]						)
      __lowerCAmelCase :						Optional[Any] = torch.from_numpy(lowercase__						).unsqueeze(0						)
      return image







class   __lowercase				(lowerCAmelCase__			):
 def __init__(							self				,							A_				,							A_				,							A_				,							) ->Optional[Any]:



       '''simple docstring'''





       super().__init__()

       self.register_modules(
           unet=_UpperCAmelCase				,							scheduler=_UpperCAmelCase				,							movq=_UpperCAmelCase				,							)
       __lowerCAmelCase :						List[str] = 2 ** (len(self.movq.config.block_out_channels ) - 1)
 def      UpperCamelCase__  (							self				,							A_				,							A_				,							A_ ) ->List[str]:



       '''simple docstring'''





       __lowerCAmelCase :						Optional[int] = min(int(num_inference_steps * strength )				,							_UpperCAmelCase )

       __lowerCAmelCase :						Optional[int] = max(num_inference_steps - init_timestep				,							0 )
       __lowerCAmelCase :						Tuple = self.scheduler.timesteps[t_start:]

       return timesteps, num_inference_steps - t_start
 def      UpperCamelCase__  (							self				,							A_				,							A_				,							A_				,							A_				,							A_				,							A_				,							A_=None ) ->Optional[Any]:



       '''simple docstring'''





       if not isinstance(_UpperCAmelCase				,							(torch.Tensor, PIL.Image.Image, list) ):
             raise ValueError(
                 f"""`image` has to be of type `torch.Tensor`, `PIL.Image.Image` or list but is {type(_UpperCAmelCase )}""" )

       __lowerCAmelCase :						Any = image.to(device=_UpperCAmelCase				,							dtype=_UpperCAmelCase )

       __lowerCAmelCase :						Tuple = batch_size * num_images_per_prompt

       if image.shape[1] == 4:
             __lowerCAmelCase :						Optional[int] = image

       else:
             if isinstance(_UpperCAmelCase				,							_UpperCAmelCase ) and len(_UpperCAmelCase ) != batch_size:
                   raise ValueError(
                       f"""You have passed a list of generators of length {len(_UpperCAmelCase )}, but requested an effective batch"""
                       f""" size of {batch_size}. Make sure the batch size matches the length of the generators.""" )

             elif isinstance(_UpperCAmelCase				,							_UpperCAmelCase ):
                   __lowerCAmelCase :						Optional[Any] = [
                       self.movq.encode(image[i : i + 1] ).latent_dist.sample(generator[i] ) for i in range(_UpperCAmelCase )
                   ]
                   __lowerCAmelCase :						str = torch.cat(_UpperCAmelCase				,							dim=0 )
             else:
                   __lowerCAmelCase :						Optional[int] = self.movq.encode(_UpperCAmelCase ).latent_dist.sample(_UpperCAmelCase )

             __lowerCAmelCase :						Any = self.movq.config.scaling_factor * init_latents

       __lowerCAmelCase :						Tuple = torch.cat([init_latents]				,							dim=0 )

       __lowerCAmelCase :						Tuple = init_latents.shape
       __lowerCAmelCase :						Tuple = randn_tensor(_UpperCAmelCase				,							generator=_UpperCAmelCase				,							device=_UpperCAmelCase				,							dtype=_UpperCAmelCase )

       # get latents
       __lowerCAmelCase :						Optional[int] = self.scheduler.add_noise(_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase )

       __lowerCAmelCase :						Union[str, Any] = init_latents

       return latents
 def      UpperCamelCase__  (							self				,							A_=0 ) ->List[str]:



       '''simple docstring'''





       if is_accelerate_available():
             from accelerate import cpu_offload
       else:
             raise ImportError('''Please install accelerate via `pip install accelerate`''' )

       __lowerCAmelCase :						str = torch.device(f"""cuda:{gpu_id}""" )

       __lowerCAmelCase :						str = [
           self.unet,
           self.movq,
       ]
       for cpu_offloaded_model in models:
             if cpu_offloaded_model is not None:
                   cpu_offload(_UpperCAmelCase				,							_UpperCAmelCase )
 def      UpperCamelCase__  (							self				,							A_=0 ) ->List[str]:



       '''simple docstring'''





       if is_accelerate_available() and is_accelerate_version('''>='''				,							'''0.17.0.dev0''' ):
             from accelerate import cpu_offload_with_hook
       else:
             raise ImportError('''`enable_model_cpu_offload` requires `accelerate v0.17.0` or higher.''' )

       __lowerCAmelCase :						Optional[Any] = torch.device(f"""cuda:{gpu_id}""" )

       if self.device.type != "cpu":
             self.to('''cpu'''				,							silence_dtype_warnings=_UpperCAmelCase )
             torch.cuda.empty_cache()  # otherwise we don't see the memory savings (but they probably exist)

       __lowerCAmelCase :						int = None
       for cpu_offloaded_model in [self.unet, self.movq]:
             __lowerCAmelCase,							__lowerCAmelCase :						Dict = cpu_offload_with_hook(_UpperCAmelCase				,							_UpperCAmelCase				,							prev_module_hook=_UpperCAmelCase )

       # We'll offload the last model manually.
       __lowerCAmelCase :						Any = hook
 @property
 # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline._execution_device
 def      UpperCamelCase__  (							self ) ->Optional[int]:



       '''simple docstring'''





       if not hasattr(self.unet				,							'''_hf_hook''' ):
             return self.device
       for module in self.unet.modules():
             if (
                 hasattr(_UpperCAmelCase				,							'''_hf_hook''' )
                 and hasattr(module._hf_hook				,							'''execution_device''' )
                 and module._hf_hook.execution_device is not None
             ):
                   return torch.device(module._hf_hook.execution_device )
       return self.device





 @torch.no_grad()
 @replace_example_docstring(_UpperCAmelCase )
 def __call__(							self				,							A_				,							A_				,							A_				,							A_				,							A_ = 512				,							A_ = 512				,							A_ = 100				,							A_ = 4.0				,							A_ = 0.3				,							A_ = 1				,							A_ = None				,							A_ = "pil"				,							A_ = True				,							) ->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase :						Any = self._execution_device

       __lowerCAmelCase :						Union[str, Any] = guidance_scale > 1.0

       if isinstance(_UpperCAmelCase				,							_UpperCAmelCase ):
             __lowerCAmelCase :						int = torch.cat(_UpperCAmelCase				,							dim=0 )
       if isinstance(_UpperCAmelCase				,							_UpperCAmelCase ):
             __lowerCAmelCase :						List[str] = torch.cat(_UpperCAmelCase				,							dim=0 )
       if isinstance(_UpperCAmelCase				,							_UpperCAmelCase ):
             __lowerCAmelCase :						Union[str, Any] = torch.cat(_UpperCAmelCase				,							dim=0 )

       __lowerCAmelCase :						List[str] = image_embeds.shape[0]

       if do_classifier_free_guidance:
             __lowerCAmelCase :						Tuple = image_embeds.repeat_interleave(_UpperCAmelCase				,							dim=0 )
             __lowerCAmelCase :						List[Any] = negative_image_embeds.repeat_interleave(_UpperCAmelCase				,							dim=0 )
             __lowerCAmelCase :						List[str] = hint.repeat_interleave(_UpperCAmelCase				,							dim=0 )

       __lowerCAmelCase :						str = torch.cat([negative_image_embeds, image_embeds]				,							dim=0 ).to(dtype=self.unet.dtype				,							device=_UpperCAmelCase )
       __lowerCAmelCase :						List[str] = torch.cat([hint, hint]				,							dim=0 ).to(dtype=self.unet.dtype				,							device=_UpperCAmelCase )

       if not isinstance(_UpperCAmelCase				,							_UpperCAmelCase ):
             __lowerCAmelCase :						Union[str, Any] = [image]
       if not all(isinstance(_UpperCAmelCase				,							(PIL.Image.Image, torch.Tensor) ) for i in image ):
             raise ValueError(
                 f"""Input is in incorrect format: {[type(_UpperCAmelCase ) for i in image]}. Currently, we only support  PIL image and pytorch tensor""" )

       __lowerCAmelCase :						List[str] = torch.cat([prepare_image(_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase ) for i in image]				,							dim=0 )
       __lowerCAmelCase :						int = image.to(dtype=image_embeds.dtype				,							device=_UpperCAmelCase )

       __lowerCAmelCase :						Tuple = self.movq.encode(_UpperCAmelCase )['''latents''']
       __lowerCAmelCase :						Union[str, Any] = latents.repeat_interleave(_UpperCAmelCase				,							dim=0 )
       self.scheduler.set_timesteps(_UpperCAmelCase				,							device=_UpperCAmelCase )
       __lowerCAmelCase,							__lowerCAmelCase :						List[Any] = self.get_timesteps(_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase )
       __lowerCAmelCase :						int = timesteps[:1].repeat(batch_size * num_images_per_prompt )
       __lowerCAmelCase,							__lowerCAmelCase :						Optional[Any] = downscale_height_and_width(_UpperCAmelCase				,							_UpperCAmelCase				,							self.movq_scale_factor )
       __lowerCAmelCase :						List[str] = self.prepare_latents(
           _UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase				,							image_embeds.dtype				,							_UpperCAmelCase				,							_UpperCAmelCase )
       for i, t in enumerate(self.progress_bar(_UpperCAmelCase ) ):
             # expand the latents if we are doing classifier free guidance
             __lowerCAmelCase :						List[Any] = torch.cat([latents] * 2 ) if do_classifier_free_guidance else latents

             __lowerCAmelCase :						int = {'''image_embeds''': image_embeds, '''hint''': hint}
             __lowerCAmelCase :						str = self.unet(
                 sample=_UpperCAmelCase				,							timestep=_UpperCAmelCase				,							encoder_hidden_states=_UpperCAmelCase				,							added_cond_kwargs=_UpperCAmelCase				,							return_dict=_UpperCAmelCase				,							)[0]

             if do_classifier_free_guidance:
                   __lowerCAmelCase,							__lowerCAmelCase :						Optional[Any] = noise_pred.split(latents.shape[1]				,							dim=1 )
                   __lowerCAmelCase,							__lowerCAmelCase :						List[str] = noise_pred.chunk(2 )
                   __lowerCAmelCase,							__lowerCAmelCase :						Union[str, Any] = variance_pred.chunk(2 )
                   __lowerCAmelCase :						Tuple = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond)
                   __lowerCAmelCase :						Tuple = torch.cat([noise_pred, variance_pred_text]				,							dim=1 )

             if not (
                 hasattr(self.scheduler.config				,							'''variance_type''' )
                 and self.scheduler.config.variance_type in ["learned", "learned_range"]
             ):
                   __lowerCAmelCase,							__lowerCAmelCase :						Tuple = noise_pred.split(latents.shape[1]				,							dim=1 )

             # compute the previous noisy sample x_t -> x_t-1

             __lowerCAmelCase :						Optional[Any] = self.scheduler.step(
                 _UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase				,							generator=_UpperCAmelCase				,							)[0]

       # post-processing
       __lowerCAmelCase :						Optional[Any] = self.movq.decode(_UpperCAmelCase				,							force_not_quantize=_UpperCAmelCase )['''sample''']

       if output_type not in ["pt", "np", "pil"]:
             raise ValueError(f"""Only the output types `pt`, `pil` and `np` are supported not output_type={output_type}""" )

       if output_type in ["np", "pil"]:
             __lowerCAmelCase :						Dict = image * 0.5 + 0.5
             __lowerCAmelCase :						int = image.clamp(0				,							1 )
             __lowerCAmelCase :						Optional[Any] = image.cpu().permute(0				,							2				,							3				,							1 ).float().numpy()

       if output_type == "pil":
             __lowerCAmelCase :						Tuple = self.numpy_to_pil(_UpperCAmelCase )

       if not return_dict:
             return (image,)

       return ImagePipelineOutput(images=_UpperCAmelCase )





import json
import os
from functools import lru_cache
from typing import TYPE_CHECKING, List, Optional, Tuple, Union

import numpy as np
import regex as re

from ...utils import is_tf_available, is_torch_available, logging


if TYPE_CHECKING:
 if is_torch_available():
  import torch
 if is_tf_available():
  import tensorflow as tf

from ...tokenization_utils import AddedToken, PreTrainedTokenizer


_UpperCamelCase			      =		logging.get_logger(__name__)

_UpperCamelCase			      =		{
    "vocab_file": "vocab.json",
    "merges_file": "merges.txt",
}

_UpperCamelCase			      =		{
    "vocab_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/vocab.json",
    },
    "merges_file": {
        "Salesforce/codegen-350M-mono": "https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/merges.txt",
    },
}

_UpperCamelCase			      =		{
    "Salesforce/codegen-350M-mono": 2048,
}







@lru_cache()
def 		_lowercase			(					):
      __lowerCAmelCase :						Union[str, Any] = (
          list(range(ord('''!'''						)       ,     ord('''~'''						) + 1						)						) + list(range(ord('''¡'''						)       ,     ord('''¬'''						) + 1						)						) + list(range(ord('''®'''						)       ,     ord('''ÿ'''						) + 1						)						)
      )
      __lowerCAmelCase :						Optional[int] = bs[:]
      __lowerCAmelCase :						Union[str, Any] = 0
      for b in range(2**8						):
            if b not in bs:
                  bs.append(lowercase__						)
                  cs.append(2**8 + n						)
                  n += 1
      __lowerCAmelCase :						Optional[int] = [chr(lowercase__						) for n in cs]
      return dict(zip(lowercase__       ,     lowercase__						)						)







def 		_lowercase			(					lowercase__						):
      __lowerCAmelCase :						List[Any] = set()
      __lowerCAmelCase :						int = word[0]
      for char in word[1:]:
            pairs.add((prev_char, char)						)
            __lowerCAmelCase :						int = char
      return pairs







class   __lowercase				(_UpperCAmelCase			):
 _UpperCamelCase        =  VOCAB_FILES_NAMES
 _UpperCamelCase        =  PRETRAINED_VOCAB_FILES_MAP
 _UpperCamelCase        =  PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
 _UpperCamelCase        =  ["""input_ids""", """attention_mask"""]
 def __init__(							self				,							A_				,							A_				,							A_="replace"				,							A_="<|endoftext|>"				,							A_="<|endoftext|>"				,							A_="<|endoftext|>"				,							A_=None				,							A_=False				,							A_=False				,							**A_				,							) ->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase :						Union[str, Any] = AddedToken(A_				,							lstrip=A_				,							rstrip=A_ ) if isinstance(A_				,							A_ ) else bos_token
       __lowerCAmelCase :						Tuple = AddedToken(A_				,							lstrip=A_				,							rstrip=A_ ) if isinstance(A_				,							A_ ) else eos_token
       __lowerCAmelCase :						Any = AddedToken(A_				,							lstrip=A_				,							rstrip=A_ ) if isinstance(A_				,							A_ ) else unk_token
       __lowerCAmelCase :						List[Any] = AddedToken(A_				,							lstrip=A_				,							rstrip=A_ ) if isinstance(A_				,							A_ ) else pad_token
       super().__init__(
           errors=A_				,							unk_token=A_				,							bos_token=A_				,							eos_token=A_				,							pad_token=A_				,							add_prefix_space=A_				,							add_bos_token=A_				,							**A_				,							)
       __lowerCAmelCase :						List[str] = add_bos_token

       with open(A_				,							encoding='''utf-8''' ) as vocab_handle:
             __lowerCAmelCase :						Optional[int] = json.load(A_ )
       __lowerCAmelCase :						str = {v: k for k, v in self.encoder.items()}
       __lowerCAmelCase :						Optional[Any] = errors  # how to handle errors in decoding
       __lowerCAmelCase :						Optional[int] = bytes_to_unicode()
       __lowerCAmelCase :						Optional[Any] = {v: k for k, v in self.byte_encoder.items()}
       with open(A_				,							encoding='''utf-8''' ) as merges_handle:
             __lowerCAmelCase :						Optional[int] = merges_handle.read().split('''\n''' )[1:-1]
       __lowerCAmelCase :						List[Any] = [tuple(merge.split() ) for merge in bpe_merges]
       __lowerCAmelCase :						List[Any] = dict(zip(A_				,							range(len(A_ ) ) ) )
       __lowerCAmelCase :						Any = {}
       __lowerCAmelCase :						Optional[int] = add_prefix_space

       # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
       __lowerCAmelCase :						Optional[Any] = re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+''' )
 @property
 def      UpperCamelCase__  (							self ) ->Dict:



       '''simple docstring'''





       return len(self.encoder )
 def      UpperCamelCase__  (							self ) ->str:



       '''simple docstring'''





       return dict(self.encoder				,							**self.added_tokens_encoder )
 def      UpperCamelCase__  (							self				,							A_ ) ->Tuple:



       '''simple docstring'''





       if token in self.cache:
             return self.cache[token]
       __lowerCAmelCase :						Tuple = tuple(A_ )
       __lowerCAmelCase :						str = get_pairs(A_ )

       if not pairs:
             return token

       while True:
             __lowerCAmelCase :						int = min(A_				,							key=lambda A_ : self.bpe_ranks.get(A_				,							float('''inf''' ) ) )
             if bigram not in self.bpe_ranks:
                   break
             __lowerCAmelCase,							__lowerCAmelCase :						Dict = bigram
             __lowerCAmelCase :						Union[str, Any] = []
             __lowerCAmelCase :						Union[str, Any] = 0
             while i < len(A_ ):
                   try:
                         __lowerCAmelCase :						List[str] = word.index(A_				,							A_ )
                   except ValueError:
                         new_word.extend(word[i:] )
                         break
                   else:
                         new_word.extend(word[i:j] )
                         __lowerCAmelCase :						Optional[int] = j

                   if word[i] == first and i < len(A_ ) - 1 and word[i + 1] == second:
                         new_word.append(first + second )
                         i += 2
                   else:
                         new_word.append(word[i] )
                         i += 1
             __lowerCAmelCase :						Optional[int] = tuple(A_ )
             __lowerCAmelCase :						List[Any] = new_word
             if len(A_ ) == 1:
                   break
             else:
                   __lowerCAmelCase :						Optional[int] = get_pairs(A_ )
       __lowerCAmelCase :						List[str] = ''' '''.join(A_ )
       __lowerCAmelCase :						Any = word
       return word
 def      UpperCamelCase__  (							self				,							A_				,							A_=None ) ->Tuple:



       '''simple docstring'''





       if self.add_bos_token:
             __lowerCAmelCase :						Any = [self.bos_token_id]
       else:
             __lowerCAmelCase :						Union[str, Any] = []

       __lowerCAmelCase :						Union[str, Any] = bos_token_ids + token_ids_a

       if token_ids_a is None:
             return output

       return output + bos_token_ids + token_ids_a
 def      UpperCamelCase__  (							self				,							A_ ) ->List[str]:



       '''simple docstring'''





       __lowerCAmelCase :						str = []
       for token in re.findall(self.pat				,							A_ ):
             __lowerCAmelCase :						Any = ''''''.join(
                 self.byte_encoder[b] for b in token.encode('''utf-8''' ) )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
             bpe_tokens.extend(bpe_token for bpe_token in self.bpe(A_ ).split(''' ''' ) )
       return bpe_tokens
 def      UpperCamelCase__  (							self				,							A_ ) ->Dict:



       '''simple docstring'''





       return self.encoder.get(A_				,							self.encoder.get(self.unk_token ) )
 def      UpperCamelCase__  (							self				,							A_ ) ->Optional[int]:



       '''simple docstring'''





       return self.decoder.get(A_ )
 def      UpperCamelCase__  (							self				,							A_ ) ->Optional[Any]:



       '''simple docstring'''





       __lowerCAmelCase :						List[Any] = ''''''.join(A_ )
       __lowerCAmelCase :						Optional[Any] = bytearray([self.byte_decoder[c] for c in text] ).decode('''utf-8'''				,							errors=self.errors )
       return text
 def      UpperCamelCase__  (							self				,							A_				,							A_ = None ) ->Tuple[str]:



       '''simple docstring'''





       if not os.path.isdir(A_ ):
             logger.error(f"""Vocabulary path ({save_directory}) should be a directory""" )
             return
       __lowerCAmelCase :						List[str] = os.path.join(
           A_				,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file'''] )
       __lowerCAmelCase :						Tuple = os.path.join(
           A_				,							(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file'''] )

       with open(A_				,							'''w'''				,							encoding='''utf-8''' ) as f:
             f.write(json.dumps(self.encoder				,							indent=2				,							sort_keys=A_				,							ensure_ascii=A_ ) + '''\n''' )

       __lowerCAmelCase :						Optional[int] = 0
       with open(A_				,							'''w'''				,							encoding='''utf-8''' ) as writer:
             writer.write('''#version: 0.2\n''' )
             for bpe_tokens, token_index in sorted(self.bpe_ranks.items()				,							key=lambda A_ : kv[1] ):
                   if index != token_index:
                         logger.warning(
                             f"""Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."""
                             ''' Please check that the tokenizer is not corrupted!''' )
                         __lowerCAmelCase :						List[Any] = token_index
                   writer.write(''' '''.join(A_ ) + '''\n''' )
                   index += 1

       return vocab_file, merge_file
 def      UpperCamelCase__  (							self				,							A_				,							A_=False				,							**A_ ) ->Union[str, Any]:



       '''simple docstring'''





       __lowerCAmelCase :						Tuple = kwargs.pop('''add_prefix_space'''				,							self.add_prefix_space )
       if is_split_into_words or add_prefix_space:
             __lowerCAmelCase :						List[str] = ''' ''' + text
       return (text, kwargs)
 def      UpperCamelCase__  (							self				,							A_				,							A_ = False				,							A_ = None				,							A_ = None				,							**A_				,							) ->str:



       '''simple docstring'''





       __lowerCAmelCase :						List[Any] = super()._decode(
           token_ids=A_				,							skip_special_tokens=A_				,							clean_up_tokenization_spaces=A_				,							**A_				,							)

       if truncate_before_pattern is not None and len(A_ ) > 0:
             __lowerCAmelCase :						str = self.truncate(A_				,							A_ )

       return decoded_text





 def      UpperCamelCase__  (							self				,							A_				,							A_ ) ->int:



       '''simple docstring'''





       def find_re(A_				,							A_				,							A_ ):
             __lowerCAmelCase :						Dict = pattern.search(A_				,							A_ )
             return m.start() if m else -1

       __lowerCAmelCase :						int = [re.compile(A_				,							re.MULTILINE ) for pattern in truncate_before_pattern]

       __lowerCAmelCase :						Tuple = list(re.finditer('''^print'''				,							A_				,							re.MULTILINE ) )

       if len(A_ ) > 1:
             __lowerCAmelCase :						int = completion[: prints[1].start()]

       __lowerCAmelCase :						Optional[Any] = list(re.finditer('''^def'''				,							A_				,							re.MULTILINE ) )

       if len(A_ ) > 1:
             __lowerCAmelCase :						List[str] = completion[: defs[1].start()]

       __lowerCAmelCase :						Optional[Any] = 0

       __lowerCAmelCase :						Optional[int] = [
           pos for pos in [find_re(A_				,							A_				,							A_ ) for terminal in terminals] if pos != -1
       ]

       if len(A_ ) > 0:
             return completion[: min(A_ )]
       else:
             return completion


