





'''simple docstring'''






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

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import is_sentencepiece_available, logging


if is_sentencepiece_available():
      from .tokenization_xglm import XGLMTokenizer
else:
      A__       :     int          =None


A__       :     List[Any]          =logging.get_logger(__name__)

A__       :     Any          ={'''vocab_file''': '''sentencepiece.bpe.model''', '''tokenizer_file''': '''tokenizer.json'''}

A__       :     Dict          ={
    '''vocab_file''': {
        '''facebook/xglm-564M''': '''https://huggingface.co/facebook/xglm-564M/resolve/main/sentencepiece.bpe.model''',
    },
    '''tokenizer_file''': {
        '''facebook/xglm-564M''': '''https://huggingface.co/facebook/xglm-564M/resolve/main/tokenizer.json''',
    },
}

A__       :     List[Any]          ={
    '''facebook/xglm-564M''': 20_48,
}




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




      def __init__(    self    :				Dict    ,    __snake_case    :				Tuple=None    ,    __snake_case    :				Optional[int]=None    ,    __snake_case    :				int="<s>"    ,    __snake_case    :				int="</s>"    ,    __snake_case    :				List[Any]="</s>"    ,    __snake_case    :				Optional[int]="<s>"    ,    __snake_case    :				str="<unk>"    ,    __snake_case    :				Union[str, Any]="<pad>"    ,    **__snake_case    :				Dict    ,    )       ->   Any:
            # Compatibility with the original tokenizer
            _lowerCAmelCase							  =							7
            _lowerCAmelCase							  =							[f"<madeupword{i}>" for i in range(self.num_madeup_words						)]

            _lowerCAmelCase							  =							kwargs.get("""additional_special_tokens"""    ,    []						)
            kwargs["additional_special_tokens"] += [
                word for word in madeup_words if word not in kwargs["additional_special_tokens"]
            ]

            super().__init__(
                __snake_case    ,    tokenizer_file=__snake_case    ,    bos_token=__snake_case    ,    eos_token=__snake_case    ,    sep_token=__snake_case    ,    cls_token=__snake_case    ,    unk_token=__snake_case    ,    pad_token=__snake_case    ,    **__snake_case    ,    )

            _lowerCAmelCase							  =							vocab_file
            _lowerCAmelCase							  =							False if not self.vocab_file else True




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

            if token_ids_a is None:
                  return [self.sep_token_id] + token_ids_a
            _lowerCAmelCase							  =							[self.sep_token_id]
            return sep + token_ids_a + sep + sep + token_ids_a




      def        lowercase__					(    self    :				str    ,    __snake_case    :				List[int]    ,    __snake_case    :				Optional[List[int]] = None						)       ->   List[int]:
            _lowerCAmelCase							  =							[self.sep_token_id]

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




      def        lowercase__					(    self    :				List[str]    ,    __snake_case    :				str    ,    __snake_case    :				Optional[str] = None						)       ->   Tuple[str]:
            if not self.can_save_slow_tokenizer:
                  raise ValueError(
                      """Your fast tokenizer does not have the necessary information to save the vocabulary for a slow """
                      """tokenizer."""						)

            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"""]						)

            if os.path.abspath(self.vocab_file						) != os.path.abspath(__snake_case						):
                  copyfile(self.vocab_file    ,    __snake_case						)

            return (out_vocab_file,)













'''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"""						)						)






