code: 
  
  
  
  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      = get_logger(__name__)
  def 	UpperCamelCase			(     snake_case__ :					str ,  snake_case__ :					List[dict] ,  snake_case__ :					list ,  snake_case__ :					list ,  snake_case__ :					int ,  snake_case__ :					GradScaler = None ,  )			->  Optional[int]:
     # Model states
     for i, state in enumerate(snake_case__			):
        UpperCamelCase					:						Optional[int]					=			F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
        UpperCamelCase					:						Tuple					=			os.path.join(snake_case__ ,  snake_case__			)
        save(snake_case__ ,  snake_case__			)
        logger.info(F"""Model weights saved in {output_model_file}"""			)
     # Optimizer states
     for i, opt in enumerate(snake_case__			):
        UpperCamelCase					:						str					=			opt.state_dict()
        UpperCamelCase					:						Any					=			F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
        UpperCamelCase					:						int					=			os.path.join(snake_case__ ,  snake_case__			)
        save(snake_case__ ,  snake_case__			)
        logger.info(F"""Optimizer state saved in {output_optimizer_file}"""			)
     # Scheduler states
     for i, scheduler in enumerate(snake_case__			):
        UpperCamelCase					:						str					=			scheduler.state_dict()
        UpperCamelCase					:						int					=			F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
        UpperCamelCase					:						str					=			os.path.join(snake_case__ ,  snake_case__			)
        save(snake_case__ ,  snake_case__			)
        logger.info(F"""Scheduler state saved in {output_scheduler_file}"""			)
     # GradScaler state
     if scaler is not None:
        UpperCamelCase					:						List[str]					=			scaler.state_dict()
        UpperCamelCase					:						List[str]					=			os.path.join(snake_case__ ,  snake_case__			)
        torch.save(snake_case__ ,  snake_case__			)
        logger.info(F"""Gradient scaler state saved in {output_scaler_file}"""			)
     # Random number generator states
     UpperCamelCase					:						Optional[Any]					=			{}
     UpperCamelCase					:						List[Any]					=			F"""{RNG_STATE_NAME}_{process_index}.pkl"""
     UpperCamelCase					:						int					=			random.getstate()
     UpperCamelCase					:						Any					=			np.random.get_state()
     UpperCamelCase					:						Dict					=			torch.get_rng_state()
     if is_xpu_available():
        UpperCamelCase					:						Any					=			torch.xpu.get_rng_state_all()
     else:
        UpperCamelCase					:						str					=			torch.cuda.get_rng_state_all()
     if is_tpu_available():
        UpperCamelCase					:						List[str]					=			xm.get_rng_state()
     UpperCamelCase					:						Tuple					=			os.path.join(snake_case__ ,  snake_case__			)
     torch.save(snake_case__ ,  snake_case__			)
     logger.info(F"""Random states saved in {output_states_file}"""			)
     return output_dir
  def 	UpperCamelCase			(     snake_case__ :					List[Any] ,  snake_case__ :					Union[str, Any] ,  snake_case__ :					List[str] ,  snake_case__ :					str ,  snake_case__ :					Dict ,  snake_case__ :					List[Any]=None ,  snake_case__ :					Tuple=None ,  **snake_case__ :					str ,  )			->  List[str]:
     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					:						List[str]					=			'cpu'
     elif map_location == "on_device":
        UpperCamelCase					:						Dict					=			PartialState().device
     # Model states
     for i, model in enumerate(snake_case__			):
        UpperCamelCase					:						Dict					=			F"""{MODEL_NAME}.bin""" if i == 0 else F"""{MODEL_NAME}_{i}.bin"""
        UpperCamelCase					:						Optional[Any]					=			os.path.join(snake_case__ ,  snake_case__			)
        models[i].load_state_dict(torch.load(snake_case__ ,  map_location=snake_case__			) ,  **snake_case__			)
     logger.info('All model weights loaded successfully'			)
  
     # Optimizer states
     for i, opt in enumerate(snake_case__			):
        UpperCamelCase					:						List[str]					=			F"""{OPTIMIZER_NAME}.bin""" if i == 0 else F"""{OPTIMIZER_NAME}_{i}.bin"""
        UpperCamelCase					:						int					=			os.path.join(snake_case__ ,  snake_case__			)
        UpperCamelCase					:						Union[str, Any]					=			torch.load(snake_case__ ,  map_location=snake_case__			)
        optimizers[i].load_state_dict(snake_case__			)
     logger.info('All optimizer states loaded successfully'			)
  
     # Scheduler states
     for i, scheduler in enumerate(snake_case__			):
        UpperCamelCase					:						Any					=			F"""{SCHEDULER_NAME}.bin""" if i == 0 else F"""{SCHEDULER_NAME}_{i}.bin"""
        UpperCamelCase					:						Tuple					=			os.path.join(snake_case__ ,  snake_case__			)
        scheduler.load_state_dict(torch.load(snake_case__			)			)
     logger.info('All scheduler states loaded successfully'			)
  
     # GradScaler state
     if scaler is not None:
        UpperCamelCase					:						int					=			os.path.join(snake_case__ ,  snake_case__			)
        scaler.load_state_dict(torch.load(snake_case__			)			)
        logger.info('GradScaler state loaded successfully'			)
  
     # Random states
     try:
        UpperCamelCase					:						Union[str, Any]					=			torch.load(os.path.join(snake_case__ ,  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			(     snake_case__ :					Optional[Any] ,  snake_case__ :					Optional[Any] ,  snake_case__ :					int = 0			)			->  Union[str, Any]:
     UpperCamelCase					:						Dict					=			Path(snake_case__			) / F"""custom_checkpoint_{index}.pkl"""
     logger.info(F"""Saving the state of {get_pretty_name(snake_case__			)} to {save_location}"""			)
     torch.save(obj.state_dict() ,  snake_case__			)
  
  
  def 	UpperCamelCase			(     snake_case__ :					List[Any] ,  snake_case__ :					int ,  snake_case__ :					int = 0			)			->  Optional[int]:
     UpperCamelCase					:						Optional[Any]					=			F"""{path}/custom_checkpoint_{index}.pkl"""
     logger.info(F"""Loading the state of {get_pretty_name(snake_case__			)} from {load_location}"""			)
     obj.load_state_dict(torch.load(snake_case__ ,  map_location='cpu'			)			)
  
  
  
  
style_context: 
  
  
  
  from typing import Optional, Tuple
  
  from ...tokenization_utils_fast import PreTrainedTokenizerFast
  from ...utils import logging
  from .tokenization_openai import OpenAIGPTTokenizer
  
  
  __UpperCAmelCase      = logging.get_logger(__name__)
  
  __UpperCAmelCase      = {'''vocab_file''': '''vocab.json''', '''merges_file''': '''merges.txt''', '''tokenizer_file''': '''tokenizer.json'''}
  
  __UpperCAmelCase      = {
      '''vocab_file''': {'''openai-gpt''': '''https://huggingface.co/openai-gpt/resolve/main/vocab.json'''},
      '''merges_file''': {'''openai-gpt''': '''https://huggingface.co/openai-gpt/resolve/main/merges.txt'''},
      '''tokenizer_file''': {'''openai-gpt''': '''https://huggingface.co/openai-gpt/resolve/main/tokenizer.json'''},
  }
  
  __UpperCAmelCase      = {
      '''openai-gpt''': 512,
  }
  
  
  
  class        lowerCAmelCase_	(						a__							):
         UpperCAmelCase__       :      Union[str, Any]       =						VOCAB_FILES_NAMES
         UpperCAmelCase__       :      int       =						PRETRAINED_VOCAB_FILES_MAP
         UpperCAmelCase__       :      Optional[int]       =						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
         UpperCAmelCase__       :      Optional[int]       =						["input_ids", "attention_mask"]
         UpperCAmelCase__       :      Tuple       =						OpenAIGPTTokenizer
         def __init__(    self,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_=None,    SCREAMING_SNAKE_CASE_="<unk>",    **SCREAMING_SNAKE_CASE_ )  ->	List[str]:
            super().__init__(SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_,    tokenizer_file=SCREAMING_SNAKE_CASE_,    unk_token=SCREAMING_SNAKE_CASE_,    **SCREAMING_SNAKE_CASE_ )
         @property
         def       snake_case_ (    self )  ->	Dict:
            return True
  
  
  
  
  
  
         def       snake_case_ (    self,    SCREAMING_SNAKE_CASE_,    SCREAMING_SNAKE_CASE_ = None )  ->	Tuple[str]:
            UpperCamelCase					:						List[Any]					=			self._tokenizer.model.save(SCREAMING_SNAKE_CASE_,    name=SCREAMING_SNAKE_CASE_ )
            return tuple(SCREAMING_SNAKE_CASE_ )
  
  
  
  