





'''simple docstring'''






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


A__       :     int          =logging.get_logger(__name__)

A__       :     List[str]          ={
    '''vocab_file''': '''vocab.json''',
    '''merges_file''': '''merges.txt''',
}

A__       :     List[Any]          ={
    '''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''',
    },
}

A__       :     List[str]          ={
    '''Salesforce/codegen-350M-mono''': 20_48,
}
@lru_cache()
def      UpperCamelCase__   (	):
      _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(__a						)
                  cs.append(2**8 + n						)
                  n += 1
      _lowerCAmelCase							  =							[chr(__a						) for n in cs]
      return dict(zip(__a							,			__a						)						)
def      UpperCamelCase__   (	lowerCAmelCase						):
      _lowerCAmelCase							  =							set()
      _lowerCAmelCase							  =							word[0]
      for char in word[1:]:
            pairs.add((prev_char, char)						)
            _lowerCAmelCase							  =							char
      return pairs





class 				UpperCAmelCase  (    UpperCamelCase__     ):
      _lowercase:    Optional[int]		  =     VOCAB_FILES_NAMES
      _lowercase:    List[str]		  =     PRETRAINED_VOCAB_FILES_MAP
      _lowercase:    str		  =     PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
      _lowercase:    int		  =     ['''input_ids''', '''attention_mask''']




      def __init__(    self    :				Any    ,    __snake_case    :				str    ,    __snake_case    :				str    ,    __snake_case    :				List[str]="replace"    ,    __snake_case    :				Union[str, Any]="<|endoftext|>"    ,    __snake_case    :				Optional[int]="<|endoftext|>"    ,    __snake_case    :				Tuple="<|endoftext|>"    ,    __snake_case    :				Tuple=None    ,    __snake_case    :				int=False    ,    __snake_case    :				str=False    ,    **__snake_case    :				Dict    ,    )       ->   Optional[Any]:
            _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        lowercase__					(    self    :				str						)       ->   Any:
            return len(self.encoder						)




      def        lowercase__					(    self    :				Dict						)       ->   Tuple:
            return dict(self.encoder    ,    **self.added_tokens_encoder						)




      def        lowercase__					(    self    :				int    ,    __snake_case    :				Tuple						)       ->   Optional[int]:
            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 __snake_case						: 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        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				str    ,    __snake_case    :				Union[str, Any]=None						)       ->   Any:
            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        lowercase__					(    self    :				Dict    ,    __snake_case    :				List[Any]						)       ->   List[Any]:
            _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        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				Optional[int]						)       ->   Dict:
            return self.encoder.get(snake_case_    ,    self.encoder.get(self.unk_token						)						)




      def        lowercase__					(    self    :				Dict    ,    __snake_case    :				Tuple						)       ->   Any:
            return self.decoder.get(snake_case_						)




      def        lowercase__					(    self    :				str    ,    __snake_case    :				int						)       ->   Any:
            _lowerCAmelCase							  =							''''''.join(snake_case_						)
            _lowerCAmelCase							  =							bytearray([self.byte_decoder[c] for c in text]						).decode("""utf-8"""    ,    errors=self.errors						)
            return text




      def        lowercase__					(    self    :				int    ,    __snake_case    :				Any    ,    __snake_case    :				List[Any] = None						)       ->   Tuple[str]:
            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 __snake_case						: 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        lowercase__					(    self    :				Any    ,    __snake_case    :				Tuple    ,    __snake_case    :				List[Any]=False    ,    **__snake_case    :				Tuple						)       ->   Dict:
            _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        lowercase__					(    self    :				Optional[int]    ,    __snake_case    :				List[str]    ,    __snake_case    :				Union[str, Any] = False    ,    __snake_case    :				Optional[Any] = None    ,    __snake_case    :				List[str] = None    ,    **__snake_case    :				int    ,    )       ->   str:
            _lowerCAmelCase							  =							super()._decode(
                token_ids=snake_case_    ,    skip_special_tokens=snake_case_    ,    clean_up_tokenization_spaces=snake_case_    ,    **snake_case_    ,    )

            if truncate_before_pattern is not None and len(snake_case_						) > 0:
                  _lowerCAmelCase							  =							self.truncate(snake_case_    ,    snake_case_						)

            return decoded_text




      def        lowercase__					(    self    :				Tuple    ,    __snake_case    :				Any    ,    __snake_case    :				int						)       ->   Union[str, Any]:
            def find_re(__snake_case    :				List[Any]    ,    __snake_case    :				Dict    ,    __snake_case    :				str						):
                  _lowerCAmelCase							  =							pattern.search(snake_case_    ,    snake_case_						)
                  return m.start() if m else -1

            _lowerCAmelCase							  =							[re.compile(snake_case_    ,    re.MULTILINE						) for pattern in truncate_before_pattern]

            _lowerCAmelCase							  =							list(re.finditer("""^print"""    ,    snake_case_    ,    re.MULTILINE						)						)

            if len(snake_case_						) > 1:
                  _lowerCAmelCase							  =							completion[: prints[1].start()]

            _lowerCAmelCase							  =							list(re.finditer("""^def"""    ,    snake_case_    ,    re.MULTILINE						)						)

            if len(snake_case_						) > 1:
                  _lowerCAmelCase							  =							completion[: defs[1].start()]

            _lowerCAmelCase							  =							0

            _lowerCAmelCase							  =							[
                pos for pos in [find_re(snake_case_    ,    snake_case_    ,    snake_case_						) for terminal in terminals] if pos != -1
            ]

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













'''simple docstring'''






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


A__       :     Optional[Any]          =get_logger(__name__)
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = None							,			):





      """simple docstring"""


      for i, state in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Model weights saved in {output_model_file}"						)
      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							opt.state_dict()
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Optimizer state saved in {output_optimizer_file}"						)
      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							scheduler.state_dict()
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Scheduler state saved in {output_scheduler_file}"						)
      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							scaler.state_dict()
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            torch.save(lowerCAmelCase							,			lowerCAmelCase						)
            logger.info(f"Gradient scaler state saved in {output_scaler_file}"						)
      # Random number generator states
      _lowerCAmelCase							  =							{}
      _lowerCAmelCase							  =							f"{RNG_STATE_NAME}_{process_index}.pkl"
      _lowerCAmelCase							  =							random.getstate()
      _lowerCAmelCase							  =							np.random.get_state()
      _lowerCAmelCase							  =							torch.get_rng_state()
      if is_xpu_available():
            _lowerCAmelCase							  =							torch.xpu.get_rng_state_all()
      else:
            _lowerCAmelCase							  =							torch.cuda.get_rng_state_all()
      if is_tpu_available():
            _lowerCAmelCase							  =							xm.get_rng_state()
      _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
      torch.save(lowerCAmelCase							,			lowerCAmelCase						)
      logger.info(f"Random states saved in {output_states_file}"						)
      return output_dir
def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase=None							,			lowerCAmelCase=None							,			**lowerCAmelCase							,			):





      """simple docstring"""


      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:
            _lowerCAmelCase							  =							"""cpu"""
      elif map_location == "on_device":
            _lowerCAmelCase							  =							PartialState().device
      # Model states
      for i, model in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{MODEL_NAME}.bin" if i == 0 else f"{MODEL_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            models[i].load_state_dict(torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)							,			**lowerCAmelCase						)
      logger.info("""All model weights loaded successfully"""						)

      # Optimizer states
      for i, opt in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{OPTIMIZER_NAME}.bin" if i == 0 else f"{OPTIMIZER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            _lowerCAmelCase							  =							torch.load(lowerCAmelCase							,			map_location=lowerCAmelCase						)
            optimizers[i].load_state_dict(lowerCAmelCase						)
      logger.info("""All optimizer states loaded successfully"""						)

      # Scheduler states
      for i, scheduler in enumerate(lowerCAmelCase						):
            _lowerCAmelCase							  =							f"{SCHEDULER_NAME}.bin" if i == 0 else f"{SCHEDULER_NAME}_{i}.bin"
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scheduler.load_state_dict(torch.load(lowerCAmelCase						)						)
      logger.info("""All scheduler states loaded successfully"""						)

      # GradScaler state
      if scaler is not None:
            _lowerCAmelCase							  =							os.path.join(lowerCAmelCase							,			lowerCAmelCase						)
            scaler.load_state_dict(torch.load(lowerCAmelCase						)						)
            logger.info("""GradScaler state loaded successfully"""						)

      # Random states
      try:
            _lowerCAmelCase							  =							torch.load(os.path.join(lowerCAmelCase							,			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      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							Path(lowerCAmelCase						) / f"custom_checkpoint_{index}.pkl"
      logger.info(f"Saving the state of {get_pretty_name(lowerCAmelCase						)} to {save_location}"						)
      torch.save(obj.state_dict()							,			lowerCAmelCase						)





def      UpperCamelCase__   (	lowerCAmelCase							,			lowerCAmelCase							,			lowerCAmelCase = 0						):





      """simple docstring"""


      _lowerCAmelCase							  =							f"{path}/custom_checkpoint_{index}.pkl"
      logger.info(f"Loading the state of {get_pretty_name(lowerCAmelCase						)} from {load_location}"						)
      obj.load_state_dict(torch.load(lowerCAmelCase							,			map_location="""cpu"""						)						)






