



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

UpperCAmelCase_          =				logging.get_logger(__name__)

UpperCAmelCase_          =				{
    'vocab_file': 'vocab.json',
    'merges_file': 'merges.txt',
}

UpperCAmelCase_          =				{
    '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',
    },
}

UpperCAmelCase_          =				{
    '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 				UpperCAmelCase__		(					)->int:
      _lowerCAmelCase			      =	(
          list(range(ord('''!'''					)  ,					ord('''~'''					) + 1					)					) + list(range(ord('''¡'''					)  ,					ord('''¬'''					) + 1					)					) + list(range(ord('''®'''					)  ,					ord('''ÿ'''					) + 1					)					)
      )
      _lowerCAmelCase			      =	bs[:]
      _lowerCAmelCase			      =	0
      for b in range(2**8					):
            if b not in bs:
                  bs.append(__lowercase					)
                  cs.append(2**8 + n					)
                  n += 1
      _lowerCAmelCase			      =	[chr(__lowercase					) for n in cs]
      return dict(zip(__lowercase  ,					__lowercase					)					)
def 				UpperCAmelCase__		(					_SCREAMING_SNAKE_CASE :						Optional[int]					)->List[str]:
      _lowerCAmelCase			      =	set()
      _lowerCAmelCase			      =	word[0]
      for char in word[1:]:
            pairs.add((prev_char, char)					)
            _lowerCAmelCase			      =	char
      return pairs





class 		UpperCAmelCase    (    snake_case__					):
   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							,  _lowerCAmelCase							,  _lowerCAmelCase							,  _lowerCAmelCase="replace"							,  _lowerCAmelCase="<|endoftext|>"							,  _lowerCAmelCase="<|endoftext|>"							,  _lowerCAmelCase="<|endoftext|>"							,  _lowerCAmelCase=None							,  _lowerCAmelCase=False							,  _lowerCAmelCase=False							,  **_lowerCAmelCase							,  ):
         _lowerCAmelCase			      =	AddedToken(__snake_case							,  lstrip=__snake_case							,  rstrip=__snake_case      ) if isinstance(__snake_case							,  __snake_case      ) else bos_token
         _lowerCAmelCase			      =	AddedToken(__snake_case							,  lstrip=__snake_case							,  rstrip=__snake_case      ) if isinstance(__snake_case							,  __snake_case      ) else eos_token
         _lowerCAmelCase			      =	AddedToken(__snake_case							,  lstrip=__snake_case							,  rstrip=__snake_case      ) if isinstance(__snake_case							,  __snake_case      ) else unk_token
         _lowerCAmelCase			      =	AddedToken(__snake_case							,  lstrip=__snake_case							,  rstrip=__snake_case      ) if isinstance(__snake_case							,  __snake_case      ) else pad_token
         super().__init__(
             errors=__snake_case							,  unk_token=__snake_case							,  bos_token=__snake_case							,  eos_token=__snake_case							,  pad_token=__snake_case							,  add_prefix_space=__snake_case							,  add_bos_token=__snake_case							,  **__snake_case							,  )
         _lowerCAmelCase			      =	add_bos_token

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

         # Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
         _lowerCAmelCase			      =	re.compile(R'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''      )





   @property
   def 						__lowerCAmelCase				(						self      ):
         return len(self.encoder      )





   def 						__lowerCAmelCase				(						self      ):
         return dict(self.encoder							,  **self.added_tokens_encoder      )





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         if token in self.cache:
               return self.cache[token]
         _lowerCAmelCase			      =	tuple(__snake_case      )
         _lowerCAmelCase			      =	get_pairs(__snake_case      )

         if not pairs:
               return token

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

                     if word[i] == first and i < len(__snake_case      ) - 1 and word[i + 1] == second:
                           new_word.append(first + second      )
                           i += 2
                     else:
                           new_word.append(word[i]      )
                           i += 1
               _lowerCAmelCase			      =	tuple(__snake_case      )
               _lowerCAmelCase			      =	new_word
               if len(__snake_case      ) == 1:
                     break
               else:
                     _lowerCAmelCase			      =	get_pairs(__snake_case      )
         _lowerCAmelCase			      =	''' '''.join(__snake_case      )
         _lowerCAmelCase			      =	word
         return word





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase=None      ):
         if self.add_bos_token:
               _lowerCAmelCase			      =	[self.bos_token_id]
         else:
               _lowerCAmelCase			      =	[]

         _lowerCAmelCase			      =	bos_token_ids + token_ids_a

         if token_ids_a is None:
               return output

         return output + bos_token_ids + token_ids_a





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase = None							,  _lowerCAmelCase = False      ):
         if already_has_special_tokens:
               return super().get_special_tokens_mask(
                   token_ids_a=__snake_case							,  token_ids_a=__snake_case							,  already_has_special_tokens=__snake_case      )

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

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





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         _lowerCAmelCase			      =	[]
         for token in re.findall(self.pat							,  __snake_case      ):
               _lowerCAmelCase			      =	''''''.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(__snake_case      ).split(''' '''      )      )
         return bpe_tokens





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         return self.encoder.get(__snake_case							,  self.encoder.get(self.unk_token      )      )





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         return self.decoder.get(__snake_case      )





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         _lowerCAmelCase			      =	''''''.join(__snake_case      )
         _lowerCAmelCase			      =	bytearray([self.byte_decoder[c] for c in text]      ).decode('''utf-8'''							,  errors=self.errors      )
         return text





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase = None      ):
         if not os.path.isdir(__snake_case      ):
               logger.error(F'''Vocabulary path ({save_directory}) should be a directory'''      )
               return
         _lowerCAmelCase			      =	os.path.join(
             __snake_case							,  (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )
         _lowerCAmelCase			      =	os.path.join(
             __snake_case							,  (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']      )

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

         _lowerCAmelCase			      =	0
         with open(__snake_case							,  '''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 _lowerCAmelCase      : 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			      =	token_index
                     writer.write(''' '''.join(__snake_case      ) + '''\n'''      )
                     index += 1

         return vocab_file, merge_file





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase							,  _lowerCAmelCase=False							,  **_lowerCAmelCase      ):
         _lowerCAmelCase			      =	kwargs.pop('''add_prefix_space'''							,  self.add_prefix_space      )
         if is_split_into_words or add_prefix_space:
               _lowerCAmelCase			      =	''' ''' + text
         return (text, kwargs)





   def 						__lowerCAmelCase				(						self							,  _lowerCAmelCase      ):
         _lowerCAmelCase			      =	[]
         for is_user, text in conversation.iter_texts():
               input_ids.extend(self.encode(__snake_case							,  add_special_tokens=__snake_case      ) + [self.eos_token_id]      )
         if len(__snake_case      ) > self.model_max_length:
               _lowerCAmelCase			      =	input_ids[-self.model_max_length :]
         return input_ids



import os
import random
from pathlib import Path
from typing import List

import numpy as np
import torch
from torch.cuda.amp import GradScaler

from .utils import (
    MODEL_NAME,
    OPTIMIZER_NAME,
    RNG_STATE_NAME,
    SCALER_NAME,
    SCHEDULER_NAME,
    get_pretty_name,
    is_tpu_available,
    is_xpu_available,
    save,
)


if is_tpu_available(check_device=False):
       import torch_xla.core.xla_model as xm

from .logging import get_logger
from .state import PartialState


__UpperCamelCase					:	int          =				get_logger(__name__)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase = None      ,							):
   # Model states
   for i, state in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Model weights saved in {output_model_file}'''							)
   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= opt.state_dict()
      UpperCAmelCase_				:   Any							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Optimizer state saved in {output_optimizer_file}'''							)
   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Tuple							= scheduler.state_dict()
      UpperCAmelCase_				:   Tuple							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   Tuple							= os.path.join(__lowercase      ,							__lowercase							)
      save(__lowercase      ,							__lowercase							)
      logger.info(F'''Scheduler state saved in {output_scheduler_file}'''							)
   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   Union[str, Any]							= scaler.state_dict()
      UpperCAmelCase_				:   Union[str, Any]							= os.path.join(__lowercase      ,							__lowercase							)
      torch.save(__lowercase      ,							__lowercase							)
      logger.info(F'''Gradient scaler state saved in {output_scaler_file}'''							)
   # Random number generator states
   UpperCAmelCase_				:   Union[str, Any]							= {}
   UpperCAmelCase_				:   Tuple							= F'''{RNG_STATE_NAME}_{process_index}.pkl'''
   UpperCAmelCase_				:   List[Any]							= random.getstate()
   UpperCAmelCase_				:   Any							= np.random.get_state()
   UpperCAmelCase_				:   List[Any]							= torch.get_rng_state()
   if is_xpu_available():
      UpperCAmelCase_				:   Optional[int]							= torch.xpu.get_rng_state_all()
   else:
      UpperCAmelCase_				:   Any							= torch.cuda.get_rng_state_all()
   if is_tpu_available():
      UpperCAmelCase_				:   List[str]							= xm.get_rng_state()
   UpperCAmelCase_				:   Any							= os.path.join(__lowercase      ,							__lowercase							)
   torch.save(__lowercase      ,							__lowercase							)
   logger.info(F'''Random states saved in {output_states_file}'''							)
   return output_dir
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase      ,							__lowercase=None      ,							__lowercase=None      ,							**__lowercase      ,							):
   if map_location not in [None, "cpu", "on_device"]:
      raise TypeError(
          '''Unsupported optimizer map location passed, please choose one of `None`, `\'cpu\'`, or `\'on_device\'`'''							)
   if map_location is None:
      UpperCAmelCase_				:   Optional[Any]							= '''cpu'''
   elif map_location == "on_device":
      UpperCAmelCase_				:   Optional[int]							= PartialState().device
   # Model states
   for i, model in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[Any]							= F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      models[i].load_state_dict(torch.load(__lowercase      ,							map_location=__lowercase							)      ,							**__lowercase							)
   logger.info('''All model weights loaded successfully'''							)

   # Optimizer states
   for i, opt in enumerate(__lowercase							):
      UpperCAmelCase_				:   List[str]							= F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   List[Any]							= os.path.join(__lowercase      ,							__lowercase							)
      UpperCAmelCase_				:   Tuple							= torch.load(__lowercase      ,							map_location=__lowercase							)
      optimizers[i].load_state_dict(__lowercase							)
   logger.info('''All optimizer states loaded successfully'''							)

   # Scheduler states
   for i, scheduler in enumerate(__lowercase							):
      UpperCAmelCase_				:   Dict							= F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
      UpperCAmelCase_				:   int							= os.path.join(__lowercase      ,							__lowercase							)
      scheduler.load_state_dict(torch.load(__lowercase							)							)
   logger.info('''All scheduler states loaded successfully'''							)

   # GradScaler state
   if scaler is not None:
      UpperCAmelCase_				:   List[str]							= os.path.join(__lowercase      ,							__lowercase							)
      scaler.load_state_dict(torch.load(__lowercase							)							)
      logger.info('''GradScaler state loaded successfully'''							)

   # Random states
   try:
      UpperCAmelCase_				:   Optional[Any]							= torch.load(os.path.join(__lowercase      ,							F'''{RNG_STATE_NAME}_{process_index}.pkl'''							)							)
      random.setstate(states['''random_state''']							)
      np.random.set_state(states['''numpy_random_seed''']							)
      torch.set_rng_state(states['''torch_manual_seed''']							)
      if is_xpu_available():
         torch.xpu.set_rng_state_all(states['''torch_xpu_manual_seed''']							)
      else:
         torch.cuda.set_rng_state_all(states['''torch_cuda_manual_seed''']							)
      if is_tpu_available():
         xm.set_rng_state(states['''xm_seed''']							)
      logger.info('''All random states loaded successfully'''							)
   except Exception:
      logger.info('''Could not load random states'''							)
def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Optional[int]							= Path(__lowercase							) / F'''custom_checkpoint_{index}.pkl'''
   logger.info(F'''Saving the state of {get_pretty_name(__lowercase							)} to {save_location}'''							)
   torch.save(obj.state_dict()      ,							__lowercase							)


def   snake_case_    (			__lowercase      ,							__lowercase      ,							__lowercase = 0							):
   UpperCAmelCase_				:   Dict							= F'''{path}/custom_checkpoint_{index}.pkl'''
   logger.info(F'''Loading the state of {get_pretty_name(__lowercase							)} from {load_location}'''							)
   obj.load_state_dict(torch.load(__lowercase      ,							map_location='''cpu'''							)							)