






"""simple docstring"""





import os

import torch

from ..logging import get_logger
from .constants import FSDP_PYTORCH_VERSION, MODEL_NAME, OPTIMIZER_NAME
from .versions import is_torch_version


if is_torch_version('''>=''', FSDP_PYTORCH_VERSION):
  import torch.distributed.checkpoint as dist_cp
  from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner, DefaultSavePlanner
  from torch.distributed.checkpoint.optimizer import load_sharded_optimizer_state_dict
  from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP
  from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType


_SCREAMING_SNAKE_CASE   :							List[Any]          =		get_logger(__name__)







def 	lowerCamelCase__			(				_lowerCamelCase			:						Any						,							_lowerCamelCase			:						List[Any]						,							_lowerCamelCase			:						Tuple						,							_lowerCamelCase			:						Union[str, Any]						,							_lowerCamelCase			:						Tuple=0					) ->							Tuple:
  os.makedirs(_lowerCamelCase						,							exist_ok=_lowerCamelCase					)
  with FSDP.state_dict_type(
      _lowerCamelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config					):
    lowerCamelCase_       	=		model.state_dict()
    if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
      lowerCamelCase_       	=		F'''{MODEL_NAME}.bin''' if model_index == 0 else F'''{MODEL_NAME}_{model_index}.bin'''
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
      if accelerator.process_index == 0:
        logger.info(F'''Saving model to {output_model_file}'''					)
        torch.save(_lowerCamelCase						,							_lowerCamelCase					)
        logger.info(F'''Model saved to {output_model_file}'''					)
    elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
      lowerCamelCase_       	=		(
          F'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
          if model_index == 0
          else F'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
      )
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
      logger.info(F'''Saving model to {output_model_file}'''					)
      torch.save(_lowerCamelCase						,							_lowerCamelCase					)
      logger.info(F'''Model saved to {output_model_file}'''					)
    elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							F'''{MODEL_NAME}_{model_index}'''					)
      os.makedirs(_lowerCamelCase						,							exist_ok=_lowerCamelCase					)
      logger.info(F'''Saving model to {ckpt_dir}'''					)
      lowerCamelCase_       	=		{'model': state_dict}

      dist_cp.save_state_dict(
          state_dict=_lowerCamelCase						,							storage_writer=dist_cp.FileSystemWriter(_lowerCamelCase					)						,							planner=DefaultSavePlanner()						,							)
      logger.info(F'''Model saved to {ckpt_dir}'''					)







def 	lowerCamelCase__			(				_lowerCamelCase			:						str						,							_lowerCamelCase			:						Dict						,							_lowerCamelCase			:						List[str]						,							_lowerCamelCase			:						str						,							_lowerCamelCase			:						Tuple=0					) ->							Optional[Any]:
  accelerator.wait_for_everyone()
  with FSDP.state_dict_type(
      _lowerCamelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config					):
    if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
      if type(_lowerCamelCase					) != FSDP and accelerator.process_index != 0:
        if not fsdp_plugin.sync_module_states:
          raise ValueError(
              'Set the `sync_module_states` flag to `True` so that model states are synced across processes when '
              'initializing FSDP object'					)
        return
      lowerCamelCase_       	=		F'''{MODEL_NAME}.bin''' if model_index == 0 else F'''{MODEL_NAME}_{model_index}.bin'''
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
      logger.info(F'''Loading model from {input_model_file}'''					)
      lowerCamelCase_       	=		torch.load(_lowerCamelCase					)
      logger.info(F'''Model loaded from {input_model_file}'''					)
    elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
      lowerCamelCase_       	=		(
          F'''{MODEL_NAME}_rank{accelerator.process_index}.bin'''
          if model_index == 0
          else F'''{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'''
      )
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
      logger.info(F'''Loading model from {input_model_file}'''					)
      lowerCamelCase_       	=		torch.load(_lowerCamelCase					)
      logger.info(F'''Model loaded from {input_model_file}'''					)
    elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
      lowerCamelCase_       	=		(
          os.path.join(_lowerCamelCase						,							F'''{MODEL_NAME}_{model_index}'''					)
          if F'''{MODEL_NAME}''' not in input_dir
          else input_dir
      )
      logger.info(F'''Loading model from {ckpt_dir}'''					)
      lowerCamelCase_       	=		{'model': model.state_dict()}
      dist_cp.load_state_dict(
          state_dict=_lowerCamelCase						,							storage_reader=dist_cp.FileSystemReader(_lowerCamelCase					)						,							planner=DefaultLoadPlanner()						,							)
      lowerCamelCase_       	=		state_dict['model']
      logger.info(F'''Model loaded from {ckpt_dir}'''					)
    model.load_state_dict(_lowerCamelCase					)







def 	lowerCamelCase__			(				_lowerCamelCase			:						Tuple						,							_lowerCamelCase			:						Tuple						,							_lowerCamelCase			:						List[str]						,							_lowerCamelCase			:						int						,							_lowerCamelCase			:						Optional[Any]						,							_lowerCamelCase			:						Union[str, Any]=0					) ->							Optional[Any]:
  os.makedirs(_lowerCamelCase						,							exist_ok=_lowerCamelCase					)
  with FSDP.state_dict_type(
      _lowerCamelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config					):
    lowerCamelCase_       	=		FSDP.optim_state_dict(_lowerCamelCase						,							_lowerCamelCase					)
    if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
      if accelerator.process_index == 0:
        lowerCamelCase_       	=		(
            F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
        )
        lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
        logger.info(F'''Saving Optimizer state to {output_optimizer_file}'''					)
        torch.save(_lowerCamelCase						,							_lowerCamelCase					)
        logger.info(F'''Optimizer state saved in {output_optimizer_file}'''					)
    else:
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							F'''{OPTIMIZER_NAME}_{optimizer_index}'''					)
      os.makedirs(_lowerCamelCase						,							exist_ok=_lowerCamelCase					)
      logger.info(F'''Saving Optimizer state to {ckpt_dir}'''					)
      dist_cp.save_state_dict(
          state_dict={'optimizer': optim_state}						,							storage_writer=dist_cp.FileSystemWriter(_lowerCamelCase					)						,							planner=DefaultSavePlanner()						,							)
      logger.info(F'''Optimizer state saved in {ckpt_dir}'''					)







def 	lowerCamelCase__			(				_lowerCamelCase			:						Tuple						,							_lowerCamelCase			:						Dict						,							_lowerCamelCase			:						List[Any]						,							_lowerCamelCase			:						int						,							_lowerCamelCase			:						List[Any]						,							_lowerCamelCase			:						Dict=0					) ->							Dict:
  accelerator.wait_for_everyone()
  with FSDP.state_dict_type(
      _lowerCamelCase						,							fsdp_plugin.state_dict_type						,							fsdp_plugin.state_dict_config						,							fsdp_plugin.optim_state_dict_config					):
    if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
      lowerCamelCase_       	=		None
      # below check should work but currently it isn't working (mostly opytorch issue),
      # in the meantime disabling it at the cost of excess memory usage
      # if accelerator.process_index == 0 or not fsdp_plugin.optim_state_dict_config.rank0_only:
      lowerCamelCase_       	=		(
          F'''{OPTIMIZER_NAME}.bin''' if optimizer_index == 0 else F'''{OPTIMIZER_NAME}_{optimizer_index}.bin'''
      )
      lowerCamelCase_       	=		os.path.join(_lowerCamelCase						,							_lowerCamelCase					)
      logger.info(F'''Loading Optimizer state from {input_optimizer_file}'''					)
      lowerCamelCase_       	=		torch.load(_lowerCamelCase					)
      logger.info(F'''Optimizer state loaded from {input_optimizer_file}'''					)
    else:
      lowerCamelCase_       	=		(
          os.path.join(_lowerCamelCase						,							F'''{OPTIMIZER_NAME}_{optimizer_index}'''					)
          if F'''{OPTIMIZER_NAME}''' not in input_dir
          else input_dir
      )
      logger.info(F'''Loading Optimizer from {ckpt_dir}'''					)
      lowerCamelCase_       	=		load_sharded_optimizer_state_dict(
          model_state_dict=model.state_dict()						,							optimizer_key='optimizer'						,							storage_reader=dist_cp.FileSystemReader(_lowerCamelCase					)						,							)
      lowerCamelCase_       	=		optim_state['optimizer']
      logger.info(F'''Optimizer loaded from {ckpt_dir}'''					)
    lowerCamelCase_       	=		FSDP.optim_state_dict_to_load(_lowerCamelCase						,							_lowerCamelCase						,							_lowerCamelCase					)
    optimizer.load_state_dict(_lowerCamelCase					)












"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


_SCREAMING_SNAKE_CASE   :							List[str]          =		{}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE   :							str          =		['''NllbTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  _SCREAMING_SNAKE_CASE   :							Optional[Any]          =		['''NllbTokenizerFast''']


if TYPE_CHECKING:
  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_nllb import NllbTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_nllb_fast import NllbTokenizerFast

else:
  import sys

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




