


'''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__  : Dict				=	get_logger(__name__)







def _UpperCamelCase	(					__A				,      __A				,      __A				,      __A				,      __A				,      __A = None				,      )							->		str:


						'''simple docstring'''



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







def _UpperCamelCase	(					__A				,      __A				,      __A				,      __A				,      __A				,      __A=None				,      __A=None				,      **__A				,      )							->		int:


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

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

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

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

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


						'''simple docstring'''



						UpperCamelCase__  =    Path(SCREAMING_SNAKE_CASE_   ) / F'''custom_checkpoint_{index}.pkl'''
						logger.info(F'''Saving the state of {get_pretty_name(SCREAMING_SNAKE_CASE_   )} to {save_location}'''   )
						torch.save(obj.state_dict()				,      SCREAMING_SNAKE_CASE_   )







def _UpperCamelCase	(					__A				,      __A				,      __A = 0   )							->		str:


						'''simple docstring'''



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









'''simple docstring'''







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
    is_vision_available,
)


a__  : Any				=	{
    'configuration_tvlt': ['TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP', 'TvltConfig'],
    'feature_extraction_tvlt': ['TvltFeatureExtractor'],
    'processing_tvlt': ['TvltProcessor'],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  a__  : int				=	[
      'TVLT_PRETRAINED_MODEL_ARCHIVE_LIST',
      'TvltModel',
      'TvltForPreTraining',
      'TvltForAudioVisualClassification',
      'TvltPreTrainedModel',
  ]

try:
  if not is_vision_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  a__  : List[Any]				=	['TvltImageProcessor']


if TYPE_CHECKING:
  from .configuration_tvlt import TVLT_PRETRAINED_CONFIG_ARCHIVE_MAP, TvltConfig
  from .processing_tvlt import TvltProcessor
  from .feature_extraction_tvlt import TvltFeatureExtractor

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tvlt import (
        TVLT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TvltForAudioVisualClassification,
        TvltForPreTraining,
        TvltModel,
        TvltPreTrainedModel,
    )

  try:
    if not is_vision_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .image_processing_tvlt import TvltImageProcessor


else:
  import sys

  a__  : Any				=	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)





