

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

import regex as re

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


if TYPE_CHECKING:
   from transformers.pipelines.conversational import Conversation

lowerCAmelCase_   =   logging.get_logger(__name__)

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

lowerCAmelCase_   =   {
    "vocab_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/vocab.json",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/vocab.json",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/vocab.json",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/vocab.json",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/vocab.json",
    },
    "merges_file": {
        "gpt2": "https://huggingface.co/gpt2/resolve/main/merges.txt",
        "gpt2-medium": "https://huggingface.co/gpt2-medium/resolve/main/merges.txt",
        "gpt2-large": "https://huggingface.co/gpt2-large/resolve/main/merges.txt",
        "gpt2-xl": "https://huggingface.co/gpt2-xl/resolve/main/merges.txt",
        "distilgpt2": "https://huggingface.co/distilgpt2/resolve/main/merges.txt",
    },
}

lowerCAmelCase_   =   {
    "gpt2": 1_0_2_4,
    "gpt2-medium": 1_0_2_4,
    "gpt2-large": 1_0_2_4,
    "gpt2-xl": 1_0_2_4,
    "distilgpt2": 1_0_2_4,
}




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




def 						A_			(  lowercase_  )			->						Dict:
       _snake_case      :							Tuple     				=						set()
       _snake_case      :							int     				=						word[0]
       for char in word[1:]:
              pairs.add((prev_char, char)  )
              _snake_case      :							str     				=						char
       return pairs

class 				A    (__UpperCAmelCase				):
      _SCREAMING_SNAKE_CASE					=							VOCAB_FILES_NAMES
      _SCREAMING_SNAKE_CASE					=							PRETRAINED_VOCAB_FILES_MAP
      _SCREAMING_SNAKE_CASE					=							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _SCREAMING_SNAKE_CASE					=							["""input_ids""", """attention_mask"""]



      def __init__(					self     ,  lowercase_     ,  lowercase_     ,  lowercase_="replace"     ,  lowercase_="<|endoftext|>"     ,  lowercase_="<|endoftext|>"     ,  lowercase_="<|endoftext|>"     ,  lowercase_=None     ,  lowercase_=False     ,  lowercase_=False     ,  **lowercase_     ,  )      ->   Optional[Any]:



             '''simple docstring'''

             _snake_case      :							Any     				=						AddedToken(lowercase_     ,  lstrip=lowercase_     ,  rstrip=lowercase_					) if isinstance(lowercase_     ,  lowercase_					) else bos_token
             _snake_case      :							int     				=						AddedToken(lowercase_     ,  lstrip=lowercase_     ,  rstrip=lowercase_					) if isinstance(lowercase_     ,  lowercase_					) else eos_token
             _snake_case      :							Optional[Any]     				=						AddedToken(lowercase_     ,  lstrip=lowercase_     ,  rstrip=lowercase_					) if isinstance(lowercase_     ,  lowercase_					) else unk_token
             _snake_case      :							str     				=						AddedToken(lowercase_     ,  lstrip=lowercase_     ,  rstrip=lowercase_					) if isinstance(lowercase_     ,  lowercase_					) else pad_token
             super().__init__(
                 errors=lowercase_     ,  unk_token=lowercase_     ,  bos_token=lowercase_     ,  eos_token=lowercase_     ,  pad_token=lowercase_     ,  add_prefix_space=lowercase_     ,  add_bos_token=lowercase_     ,  **lowercase_     ,  )
             _snake_case      :							Optional[Any]     				=						add_bos_token

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

             # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
             _snake_case      :							Union[str, 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    __a (					self					)      ->   Optional[Any]:



             '''simple docstring'''

             return len(self.encoder					)



      def    __a (					self					)      ->   Dict:



             '''simple docstring'''

             return dict(self.encoder     ,  **self.added_tokens_encoder					)



      def    __a (					self     ,  lowercase_					)      ->   Union[str, Any]:



             '''simple docstring'''

             if token in self.cache:
                    return self.cache[token]
             _snake_case      :							List[str]     				=						tuple(lowercase_					)
             _snake_case      :							Union[str, Any]     				=						get_pairs(lowercase_					)

             if not pairs:
                    return token

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

                           if word[i] == first and i < len(lowercase_					) - 1 and word[i + 1] == second:
                                  new_word.append(first + second					)
                                  i += 2
                           else:
                                  new_word.append(word[i]					)
                                  i += 1
                    _snake_case      :							Any     				=						tuple(lowercase_					)
                    _snake_case      :							Optional[int]     				=						new_word
                    if len(lowercase_					) == 1:
                           break
                    else:
                           _snake_case      :							Dict     				=						get_pairs(lowercase_					)
             _snake_case      :							int     				=						''' '''.join(lowercase_					)
             _snake_case      :							str     				=						word
             return word



      def    __a (					self     ,  lowercase_     ,  lowercase_=None					)      ->   int:



             '''simple docstring'''

             if self.add_bos_token:
                    _snake_case      :							Dict     				=						[self.bos_token_id]
             else:
                    _snake_case      :							Optional[Any]     				=						[]

             _snake_case      :							Dict     				=						bos_token_ids + token_ids_a

             if token_ids_a is None:
                    return output

             return output + bos_token_ids + token_ids_a



      def    __a (					self     ,  lowercase_     ,  lowercase_ = None     ,  lowercase_ = False					)      ->   List[int]:



             '''simple docstring'''

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

             if not self.add_bos_token:
                    return super().get_special_tokens_mask(
                        token_ids_a=lowercase_     ,  token_ids_a=lowercase_     ,  already_has_special_tokens=lowercase_					)

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



      def    __a (					self     ,  lowercase_					)      ->   int:



             '''simple docstring'''

             _snake_case      :							List[str]     				=						[]
             for token in re.findall(self.pat     ,  lowercase_					):
                    _snake_case      :							List[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(lowercase_					).split(''' '''					)					)
             return bpe_tokens



      def    __a (					self     ,  lowercase_					)      ->   Tuple:



             '''simple docstring'''

             return self.encoder.get(lowercase_     ,  self.encoder.get(self.unk_token					)					)



      def    __a (					self     ,  lowercase_					)      ->   Optional[int]:



             '''simple docstring'''

             return self.decoder.get(lowercase_					)



      def    __a (					self     ,  lowercase_					)      ->   List[str]:



             '''simple docstring'''

             _snake_case      :							int     				=						''''''.join(lowercase_					)
             _snake_case      :							List[str]     				=						bytearray([self.byte_decoder[c] for c in text]					).decode('''utf-8'''     ,  errors=self.errors					)
             return text



      def    __a (					self     ,  lowercase_     ,  lowercase_ = None					)      ->   Tuple[str]:



             '''simple docstring'''

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

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

             _snake_case      :							Optional[int]     				=						0
             with open(lowercase_     ,  '''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 lowercase_					: 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!'''					)
                                  _snake_case      :							Optional[Any]     				=						token_index
                           writer.write(''' '''.join(lowercase_					) + '''\n'''					)
                           index += 1

             return vocab_file, merge_file



      def    __a (					self     ,  lowercase_     ,  lowercase_=False     ,  **lowercase_					)      ->   Union[str, Any]:



             '''simple docstring'''

             _snake_case      :							List[str]     				=						kwargs.pop('''add_prefix_space'''     ,  self.add_prefix_space					)
             if is_split_into_words or add_prefix_space:
                    _snake_case      :							int     				=						''' ''' + text
             return (text, kwargs)



      def    __a (					self     ,  lowercase_					)      ->   List[int]:



             '''simple docstring'''

             _snake_case      :							List[str]     				=						[]
             for is_user, text in conversation.iter_texts():
                    input_ids.extend(self.encode(lowercase_     ,  add_special_tokens=lowercase_					) + [self.eos_token_id]					)
             if len(lowercase_					) > self.model_max_length:
                    _snake_case      :							Tuple     				=						input_ids[-self.model_max_length :]
             return input_ids













from functools import partial
from typing import Any, Callable, Dict, List, Type, TypeVar, Union, overload

import torch
import torch.nn as nn
import torch.types




def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						torch.Tensor:
       # The first operation in a checkpoint can't be in-place, but it's
       # nice to have in-place addition during inference. Thus...
       if not inplace:
              _snake_case      :							Any     				=						ma + ma
       else:
              ma += ma

       return ma




def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
       _snake_case      :							str     				=						-1 * len(lowercase_  )
       _snake_case      :							Any     				=						list(range(len(tensor.shape[:zero_index]  )  )  )
       return tensor.permute(first_inds + [zero_index + i for i in inds]  )




def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
       return t.reshape(t.shape[:-no_dims] + (-1,)  )




def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_      ,		lowercase_ = 1E-4  )			->						torch.Tensor:
       _snake_case      :							Any     				=						mask.expand(*value.shape  )
       return torch.sum(mask * value      ,		dim=lowercase_  ) / (eps + torch.sum(lowercase_      ,		dim=lowercase_  ))




def 						A_			(  lowercase_      ,		lowercase_ = 2.31_25      ,		lowercase_ = 21.68_75      ,		lowercase_ = 64  )			->						torch.Tensor:
       _snake_case      :							int     				=						torch.linspace(lowercase_      ,		lowercase_      ,		no_bins - 1      ,		device=pts.device  )
       _snake_case      :							Optional[Any]     				=						torch.sqrt(torch.sum((pts.unsqueeze(-2  ) - pts.unsqueeze(-3  )) ** 2      ,		dim=-1  )  )
       return torch.bucketize(lowercase_      ,		lowercase_  )




def 						A_			(  lowercase_      ,		lowercase_  )			->						dict:
       _snake_case      :							Dict     				=						dicts[0]
       _snake_case      :							str     				=						{}
       for k, v in first.items():
              _snake_case      :							List[Any]     				=						[d[k] for d in dicts]
              if isinstance(lowercase_      ,		lowercase_  ):
                     _snake_case      :							Any     				=						dict_multimap(lowercase_      ,		lowercase_  )
              else:
                     _snake_case      :							Dict     				=						fn(lowercase_  )

       return new_dict




def 						A_			(  lowercase_      ,		lowercase_  )			->						torch.Tensor:
       _snake_case      :							str     				=						v_bins.view(((1,) * len(x.shape  )) + (len(lowercase_  ),)  )
       _snake_case      :							Union[str, Any]     				=						x[..., None] - reshaped_bins
       _snake_case      :							Optional[Any]     				=						torch.argmin(torch.abs(lowercase_  )      ,		dim=-1  )
       return nn.functional.one_hot(lowercase_      ,		num_classes=len(lowercase_  )  ).float()




def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_ = 0      ,		lowercase_ = 0  )			->						torch.Tensor:
       _snake_case      :							List[Union[slice, torch.Tensor]]     				=						[]
       for i, s in enumerate(data.shape[:no_batch_dims]  ):
              _snake_case      :							Optional[Any]     				=						torch.arange(lowercase_  )
              _snake_case      :							Union[str, Any]     				=						r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape  ) - i - 1)))  )
              ranges.append(lowercase_  )

       _snake_case      :							List[Union[slice, torch.Tensor]]     				=						[slice(lowercase_  ) for _ in range(len(data.shape  ) - no_batch_dims  )]
       _snake_case      :							Union[str, Any]     				=						inds
       ranges.extend(lowercase_  )
       # Matt note: Editing this to get around the behaviour of using a list as an array index changing
       # in recent Numpy versions
       return data[tuple(lowercase_  )]


lowerCAmelCase_   =   TypeVar("T")




def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Dict[Any, Union[dict, list, tuple, Any]]:
       _snake_case      :							Dict[Any, Union[dict, list, tuple, Any]]     				=						{}
       for k, v in dic.items():
              if isinstance(lowercase_      ,		lowercase_  ):
                     _snake_case      :							Dict     				=						dict_map(lowercase_      ,		lowercase_      ,		lowercase_  )
              else:
                     _snake_case      :							List[Any]     				=						tree_map(lowercase_      ,		lowercase_      ,		lowercase_  )

       return new_dict




@overload
def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Any:
       ...




@overload
def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						dict:
       ...




@overload
def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						list:
       ...




@overload
def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						tuple:
       ...




def 						A_			(  lowercase_      ,		lowercase_      ,		lowercase_  )			->						Any:
       if isinstance(lowercase_      ,		lowercase_  ):
              return dict_map(lowercase_      ,		lowercase_      ,		lowercase_  )
       elif isinstance(lowercase_      ,		lowercase_  ):
              return [tree_map(lowercase_      ,		lowercase_      ,		lowercase_  ) for x in tree]
       elif isinstance(lowercase_      ,		lowercase_  ):
              return tuple(tree_map(lowercase_      ,		lowercase_      ,		lowercase_  ) for x in tree  )
       elif isinstance(lowercase_      ,		lowercase_  ):
              return fn(lowercase_  )
       else:
              print(type(lowercase_  )  )
              raise ValueError('''Not supported'''  )
lowerCAmelCase_   =   partial(tree_map, leaf_type=torch.Tensor)






