code: 
  
  
  
  
  '''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
  
  
  _lowerCamelCase							:		int  			=						get_logger(__name__)
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_		,							lowercase_		,							lowercase_		,							lowercase_		,							lowercase_		,							lowercase_ = None		,							):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		for i, state in enumerate(lowercase_						):
  				_snake_case			 = f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
  				_snake_case			 = 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_						):
  				_snake_case			 = opt.state_dict()
  				_snake_case			 = f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
  				_snake_case			 = 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_						):
  				_snake_case			 = scheduler.state_dict()
  				_snake_case			 = f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
  				_snake_case			 = 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:
  				_snake_case			 = scaler.state_dict()
  				_snake_case			 = os.path.join(lowercase_		,							lowercase_						)
  				torch.save(lowercase_		,							lowercase_						)
  				logger.info(f'''Gradient scaler state saved in {output_scaler_file}'''						)
  		# Random number generator states
  		_snake_case			 = {}
  		_snake_case			 = f'''{RNG_STATE_NAME}_{process_index}.pkl'''
  		_snake_case			 = random.getstate()
  		_snake_case			 = np.random.get_state()
  		_snake_case			 = torch.get_rng_state()
  		if is_xpu_available():
  				_snake_case			 = torch.xpu.get_rng_state_all()
  		else:
  				_snake_case			 = torch.cuda.get_rng_state_all()
  		if is_tpu_available():
  				_snake_case			 = xm.get_rng_state()
  		_snake_case			 = os.path.join(lowercase_		,							lowercase_						)
  		torch.save(lowercase_		,							lowercase_						)
  		logger.info(f'''Random states saved in {output_states_file}'''						)
  		return output_dir
  
  
  
  
  
  
  def 	_lowercase	(			lowercase_		,							lowercase_		,							lowercase_		,							lowercase_		,							lowercase_		,							lowercase_=None		,							lowercase_=None		,							**lowercase_		,							):
  
  
  
  		'''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:
  				_snake_case			 = "cpu"
  		elif map_location == "on_device":
  				_snake_case			 = PartialState().device
  		# Model states
  		for i, model in enumerate(lowercase_						):
  				_snake_case			 = f'''{MODEL_NAME}.bin''' if i == 0 else f'''{MODEL_NAME}_{i}.bin'''
  				_snake_case			 = 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_						):
  				_snake_case			 = f'''{OPTIMIZER_NAME}.bin''' if i == 0 else f'''{OPTIMIZER_NAME}_{i}.bin'''
  				_snake_case			 = os.path.join(lowercase_		,							lowercase_						)
  				_snake_case			 = 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_						):
  				_snake_case			 = f'''{SCHEDULER_NAME}.bin''' if i == 0 else f'''{SCHEDULER_NAME}_{i}.bin'''
  				_snake_case			 = 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:
  				_snake_case			 = os.path.join(lowercase_		,							lowercase_						)
  				scaler.load_state_dict(torch.load(lowercase_						)						)
  				logger.info("GradScaler state loaded successfully"						)
  
  		# Random states
  		try:
  				_snake_case			 = 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 	_lowercase	(			lowercase_		,							lowercase_		,							lowercase_ = 0						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = 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 	_lowercase	(			lowercase_		,							lowercase_		,							lowercase_ = 0						):
  
  
  
  		'''simple docstring'''
  
  
  
  
  
  
  
  		_snake_case			 = 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"						)						)
  
  
  
  
  
  
style_context: 
  
  
  
  
  '''simple docstring'''
  
  
  
  
  
  
  
  import unittest
  import warnings
  from dataclasses import dataclass
  
  from transformers.convert_slow_tokenizer import SpmConverter
  from transformers.testing_utils import get_tests_dir
  
  
  
  
  
  @dataclass
  class   lowerCamelCase_							:
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					lowercase__					=				42
  
  
  
  
  
  class   lowerCamelCase_							( unittest.TestCase							):
  
  
  
  
  					'''simple docstring'''
  
  
  
  
  
  
  					def        __SCREAMING_SNAKE_CASE			(					self				):
  							_snake_case			 = get_tests_dir("fixtures/test_sentencepiece.model"				)
  							_snake_case			 = get_tests_dir("fixtures/test_sentencepiece_with_bytefallback.model"				)
  
  							_snake_case			 = FakeOriginalTokenizer(vocab_file=__UpperCamelCase				)
  
  							with warnings.catch_warnings(record=__UpperCamelCase				) as w:
  									_snake_case			 = SpmConverter(__UpperCamelCase				)
  							self.assertEqual(len(__UpperCamelCase				),						0				)
  
  							_snake_case			 = FakeOriginalTokenizer(vocab_file=__UpperCamelCase				)
  
  							with warnings.catch_warnings(record=__UpperCamelCase				) as w:
  									_snake_case			 = SpmConverter(__UpperCamelCase				)
  							self.assertEqual(len(__UpperCamelCase				),						1				)
  
  							self.assertIn(
  							    "The sentencepiece tokenizer that you are converting to a fast tokenizer uses the byte fallback option"
  							    " which is not implemented in the fast tokenizers.",						str(w[0].message				),						)
  
  
  
  
  
  