




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


lowerCamelCase       :       Tuple        =    get_logger(__name__)




def 						SCREAMING_SNAKE_CASE__					(    lowercase						,lowercase						,lowercase						,lowercase						,lowercase=0     )  ->   Optional[int]:
   os.makedirs(__UpperCamelCase						,exist_ok=__UpperCamelCase     )
   with FSDP.state_dict_type(
       __UpperCamelCase						,fsdp_plugin.state_dict_type						,fsdp_plugin.state_dict_config						,fsdp_plugin.optim_state_dict_config     ):
      snake_case      :				Any				=       model.state_dict()
      if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
         snake_case      :				Tuple				=       f"""{MODEL_NAME}.bin""" if model_index == 0 else f"""{MODEL_NAME}_{model_index}.bin"""
         snake_case      :				Optional[int]				=       os.path.join(__UpperCamelCase						,__UpperCamelCase     )
         if accelerator.process_index == 0:
            logger.info(f"""Saving model to {output_model_file}"""     )
            torch.save(__UpperCamelCase						,__UpperCamelCase     )
            logger.info(f"""Model saved to {output_model_file}"""     )
      elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
         snake_case      :				Optional[Any]				=       (
             f"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
             if model_index == 0
             else f"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
         )
         snake_case      :				Optional[int]				=       os.path.join(__UpperCamelCase						,__UpperCamelCase     )
         logger.info(f"""Saving model to {output_model_file}"""     )
         torch.save(__UpperCamelCase						,__UpperCamelCase     )
         logger.info(f"""Model saved to {output_model_file}"""     )
      elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
         snake_case      :				Tuple				=       os.path.join(__UpperCamelCase						,f"""{MODEL_NAME}_{model_index}"""     )
         os.makedirs(__UpperCamelCase						,exist_ok=__UpperCamelCase     )
         logger.info(f"""Saving model to {ckpt_dir}"""     )
         snake_case      :				Any				=       {"""model""": state_dict}

         dist_cp.save_state_dict(
             state_dict=__UpperCamelCase						,storage_writer=dist_cp.FileSystemWriter(__UpperCamelCase     )						,planner=DefaultSavePlanner()						,)
         logger.info(f"""Model saved to {ckpt_dir}"""     )




def 						SCREAMING_SNAKE_CASE__					(    lowercase						,lowercase						,lowercase						,lowercase						,lowercase=0     )  ->   Optional[Any]:
   accelerator.wait_for_everyone()
   with FSDP.state_dict_type(
       __UpperCamelCase						,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(__UpperCamelCase     ) != 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
         snake_case      :				str				=       f"""{MODEL_NAME}.bin""" if model_index == 0 else f"""{MODEL_NAME}_{model_index}.bin"""
         snake_case      :				Tuple				=       os.path.join(__UpperCamelCase						,__UpperCamelCase     )
         logger.info(f"""Loading model from {input_model_file}"""     )
         snake_case      :				Tuple				=       torch.load(__UpperCamelCase     )
         logger.info(f"""Model loaded from {input_model_file}"""     )
      elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
         snake_case      :				str				=       (
             f"""{MODEL_NAME}_rank{accelerator.process_index}.bin"""
             if model_index == 0
             else f"""{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin"""
         )
         snake_case      :				Any				=       os.path.join(__UpperCamelCase						,__UpperCamelCase     )
         logger.info(f"""Loading model from {input_model_file}"""     )
         snake_case      :				List[str]				=       torch.load(__UpperCamelCase     )
         logger.info(f"""Model loaded from {input_model_file}"""     )
      elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
         snake_case      :				Tuple				=       (
             os.path.join(__UpperCamelCase						,f"""{MODEL_NAME}_{model_index}"""     )
             if f"""{MODEL_NAME}""" not in input_dir
             else input_dir
         )
         logger.info(f"""Loading model from {ckpt_dir}"""     )
         snake_case      :				Union[str, Any]				=       {"""model""": model.state_dict()}
         dist_cp.load_state_dict(
             state_dict=__UpperCamelCase						,storage_reader=dist_cp.FileSystemReader(__UpperCamelCase     )						,planner=DefaultLoadPlanner()						,)
         snake_case      :				List[str]				=       state_dict["""model"""]
         logger.info(f"""Model loaded from {ckpt_dir}"""     )
      model.load_state_dict(__UpperCamelCase     )




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




def 						SCREAMING_SNAKE_CASE__					(    lowercase						,lowercase						,lowercase						,lowercase						,lowercase						,lowercase=0     )  ->   int:
   accelerator.wait_for_everyone()
   with FSDP.state_dict_type(
       __UpperCamelCase						,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:
         snake_case      :				List[str]				=       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:
         snake_case      :				Optional[int]				=       (
             f"""{OPTIMIZER_NAME}.bin""" if optimizer_index == 0 else f"""{OPTIMIZER_NAME}_{optimizer_index}.bin"""
         )
         snake_case      :				Dict				=       os.path.join(__UpperCamelCase						,__UpperCamelCase     )
         logger.info(f"""Loading Optimizer state from {input_optimizer_file}"""     )
         snake_case      :				int				=       torch.load(__UpperCamelCase     )
         logger.info(f"""Optimizer state loaded from {input_optimizer_file}"""     )
      else:
         snake_case      :				Union[str, Any]				=       (
             os.path.join(__UpperCamelCase						,f"""{OPTIMIZER_NAME}_{optimizer_index}"""     )
             if f"""{OPTIMIZER_NAME}""" not in input_dir
             else input_dir
         )
         logger.info(f"""Loading Optimizer from {ckpt_dir}"""     )
         snake_case      :				Optional[int]				=       load_sharded_optimizer_state_dict(
             model_state_dict=model.state_dict()						,optimizer_key="""optimizer"""						,storage_reader=dist_cp.FileSystemReader(__UpperCamelCase     )						,)
         snake_case      :				Optional[Any]				=       optim_state["""optimizer"""]
         logger.info(f"""Optimizer loaded from {ckpt_dir}"""     )
      snake_case      :				Optional[Any]				=       FSDP.optim_state_dict_to_load(__UpperCamelCase						,__UpperCamelCase						,__UpperCamelCase     )
      optimizer.load_state_dict(__UpperCamelCase     )






import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin


@require_jieba
class      lowerCamelCase    (SCREAMING_SNAKE_CASE__     ,					unittest.TestCase    ):
  """simple docstring"""







  lowerCamelCase__   =  CpmAntTokenizer
  lowerCamelCase__   =  False







  def       __A   (							self       :				List[str]      )     ->		str:
        super().setUp()

        SCREAMING_SNAKE_CASE_      =     [
            "<d>",
            "</d>",
            "<s>",
            "</s>",
            "</_>",
            "<unk>",
            "<pad>",
            "</n>",
            "我",
            "是",
            "C",
            "P",
            "M",
            "A",
            "n",
            "t",
        ]
        SCREAMING_SNAKE_CASE_      =     os.path.join(self.tmpdirname     ,   VOCAB_FILES_NAMES["vocab_file"]      )
        with open(self.vocab_file     ,   "w"     ,   encoding="utf-8"      ) as vocab_writer:
              vocab_writer.write("".join([x + "\n" for x in vocab_tokens]      )      )







  @tooslow
  def       __A   (							self       :				int      )     ->		Union[str, Any]:
        SCREAMING_SNAKE_CASE_      =     CpmAntTokenizer.from_pretrained("openbmb/cpm-ant-10b"      )
        SCREAMING_SNAKE_CASE_      =     "今天天气真好！"
        SCREAMING_SNAKE_CASE_      =     ["今天", "天气", "真", "好", "！"]
        SCREAMING_SNAKE_CASE_      =     tokenizer.tokenize(__magic_name__      )
        self.assertListEqual(__magic_name__     ,   __magic_name__      )
        SCREAMING_SNAKE_CASE_      =     "今天天气真好！"
        SCREAMING_SNAKE_CASE_      =     [tokenizer.bos_token] + tokens

        SCREAMING_SNAKE_CASE_      =     [6, 9_802, 14_962, 2_082, 831, 244]
        self.assertListEqual(tokenizer.convert_tokens_to_ids(__magic_name__      )     ,   __magic_name__      )

        SCREAMING_SNAKE_CASE_      =     tokenizer.decode(__magic_name__      )
        self.assertEqual(__magic_name__     ,   __magic_name__      )


