




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


UpperCamelCase__    			=			get_logger(__name__)







def 	_a      (						SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				Optional[Any]      ,      SCREAMING_SNAKE_CASE_							:				str=0     ):
					os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															if accelerator.process_index == 0:
																				logger.info(F"""Saving model to {output_model_file}"""     )
																				torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving model to {output_model_file}"""     )
															torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model saved to {output_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{MODEL_NAME}_{model_index}"""     )
															os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving model to {ckpt_dir}"""     )
															__lowerCAmelCase				   =       {"model": state_dict}

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







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Any=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      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(SCREAMING_SNAKE_CASE_     ) != 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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading model from {input_model_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Model loaded from {input_model_file}"""     )
										elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_SNAKE_CASE_      ,      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=SCREAMING_SNAKE_CASE_      ,      storage_reader=dist_cp.FileSystemReader(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultLoadPlanner()      ,      )
															__lowerCAmelCase				   =       state_dict["model"]
															logger.info(F"""Model loaded from {ckpt_dir}"""     )
										model.load_state_dict(SCREAMING_SNAKE_CASE_     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				int      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				Optional[Any]      ,      SCREAMING_SNAKE_CASE_							:				Optional[int]      ,      SCREAMING_SNAKE_CASE_							:				str=0     ):
					os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      fsdp_plugin.state_dict_type      ,      fsdp_plugin.state_dict_config      ,      fsdp_plugin.optim_state_dict_config     ):
										__lowerCAmelCase				   =       FSDP.optim_state_dict(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Saving Optimizer state to {output_optimizer_file}"""     )
																				torch.save(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
																				logger.info(F"""Optimizer state saved in {output_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       os.path.join(SCREAMING_SNAKE_CASE_      ,      F"""{OPTIMIZER_NAME}_{optimizer_index}"""     )
															os.makedirs(SCREAMING_SNAKE_CASE_      ,      exist_ok=SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Saving Optimizer state to {ckpt_dir}"""     )
															dist_cp.save_state_dict(
															    state_dict={"optimizer": optim_state}      ,      storage_writer=dist_cp.FileSystemWriter(SCREAMING_SNAKE_CASE_     )      ,      planner=DefaultSavePlanner()      ,      )
															logger.info(F"""Optimizer state saved in {ckpt_dir}"""     )







def 	_a      (						SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Dict      ,      SCREAMING_SNAKE_CASE_							:				Tuple      ,      SCREAMING_SNAKE_CASE_							:				List[str]      ,      SCREAMING_SNAKE_CASE_							:				str      ,      SCREAMING_SNAKE_CASE_							:				Dict=0     ):
					accelerator.wait_for_everyone()
					with FSDP.state_dict_type(
					    SCREAMING_SNAKE_CASE_      ,      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(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Loading Optimizer state from {input_optimizer_file}"""     )
															__lowerCAmelCase				   =       torch.load(SCREAMING_SNAKE_CASE_     )
															logger.info(F"""Optimizer state loaded from {input_optimizer_file}"""     )
										else:
															__lowerCAmelCase				   =       (
															    os.path.join(SCREAMING_SNAKE_CASE_      ,      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(SCREAMING_SNAKE_CASE_     )      ,      )
															__lowerCAmelCase				   =       optim_state["optimizer"]
															logger.info(F"""Optimizer loaded from {ckpt_dir}"""     )
										__lowerCAmelCase				   =       FSDP.optim_state_dict_to_load(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_     )
										optimizer.load_state_dict(SCREAMING_SNAKE_CASE_     )










import math
import time
from typing import Dict, List, Optional

from torch.utils.data import Dataset

from transformers import SeqaSeqTrainer, is_torch_tpu_available
from transformers.trainer_utils import PredictionOutput, speed_metrics


if is_torch_tpu_available(check_device=False):
					import torch_xla.core.xla_model as xm
					import torch_xla.debug.metrics as met
class 			a__    (					snake_case__      ):






						def __init__(   self		,			*_A		,			_A=None		,			_A=None		,			**_A			):




											"""simple docstring"""



											super().__init__(*_A		,			**_A			)
											__lowerCAmelCase				   =       eval_examples
											__lowerCAmelCase				   =       post_process_function






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A = None		,			_A=None		,			_A = None		,			_A = "eval"		,			**_A		,			):




											"""simple docstring"""



											__lowerCAmelCase				   =       gen_kwargs.copy()
											__lowerCAmelCase				   =       (
											    gen_kwargs["max_length"] if gen_kwargs.get("max_length"			) is not None else self.args.generation_max_length
											)
											__lowerCAmelCase				   =       (
											    gen_kwargs["num_beams"] if gen_kwargs.get("num_beams"			) is not None else self.args.generation_num_beams
											)
											__lowerCAmelCase				   =       gen_kwargs

											__lowerCAmelCase				   =       self.eval_dataset if eval_dataset is None else eval_dataset
											__lowerCAmelCase				   =       self.get_eval_dataloader(_A			)
											__lowerCAmelCase				   =       self.eval_examples if eval_examples is None else eval_examples

											# Temporarily disable metric computation, we will do it in the loop here.
											__lowerCAmelCase				   =       self.compute_metrics
											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       time.time()
											__lowerCAmelCase				   =       self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__lowerCAmelCase				   =       eval_loop(
																    _A		,			description="Evaluation"		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=_A		,			metric_key_prefix=_A		,			)
											finally:
																__lowerCAmelCase				   =       compute_metrics
											__lowerCAmelCase				   =       self.args.eval_batch_size * self.args.world_size
											if f"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
																start_time += output.metrics[f"""{metric_key_prefix}_jit_compilation_time"""]
											output.metrics.update(
											    speed_metrics(
											        _A		,			_A		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size			)		,			)			)

											if self.post_process_function is not None and self.compute_metrics is not None and self.args.should_save:
																# Only the main node write the results by default
																__lowerCAmelCase				   =       self.post_process_function(_A		,			_A		,			_A			)
																__lowerCAmelCase				   =       self.compute_metrics(_A			)

																# Prefix all keys with metric_key_prefix + '_'
																for key in list(metrics.keys()			):
																					if not key.startswith(f"""{metric_key_prefix}_"""			):
																										__lowerCAmelCase				   =       metrics.pop(_A			)

																metrics.update(output.metrics			)
											else:
																__lowerCAmelCase				   =       output.metrics

											if self.args.should_log:
																# Only the main node log the results by default
																self.log(_A			)

											if self.args.tpu_metrics_debug or self.args.debug:
																# tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
																xm.master_print(met.metrics_report()			)

											__lowerCAmelCase				   =       self.callback_handler.on_evaluate(self.args		,			self.state		,			self.control		,			_A			)
											return metrics






						def 				__SCREAMING_SNAKE_CASE(   self		,			_A		,			_A		,			_A=None		,			_A = "test"		,			**_A			):




											"""simple docstring"""



											__lowerCAmelCase				   =       gen_kwargs.copy()

											__lowerCAmelCase				   =       self.get_test_dataloader(_A			)

											# Temporarily disable metric computation, we will do it in the loop here.
											__lowerCAmelCase				   =       self.compute_metrics
											__lowerCAmelCase				   =       None
											__lowerCAmelCase				   =       time.time()
											__lowerCAmelCase				   =       self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
											try:
																__lowerCAmelCase				   =       eval_loop(
																    _A		,			description="Prediction"		,			prediction_loss_only=True if compute_metrics is None else None		,			ignore_keys=_A		,			metric_key_prefix=_A		,			)
											finally:
																__lowerCAmelCase				   =       compute_metrics

											__lowerCAmelCase				   =       self.args.eval_batch_size * self.args.world_size
											if f"""{metric_key_prefix}_jit_compilation_time""" in output.metrics:
																start_time += output.metrics[f"""{metric_key_prefix}_jit_compilation_time"""]
											output.metrics.update(
											    speed_metrics(
											        _A		,			_A		,			num_samples=output.num_samples		,			num_steps=math.ceil(output.num_samples / total_batch_size			)		,			)			)
											if self.post_process_function is None or self.compute_metrics is None:
																return output

											__lowerCAmelCase				   =       self.post_process_function(_A		,			_A		,			_A		,			"predict"			)
											__lowerCAmelCase				   =       self.compute_metrics(_A			)

											# Prefix all keys with metric_key_prefix + '_'
											for key in list(metrics.keys()			):
																if not key.startswith(f"""{metric_key_prefix}_"""			):
																					__lowerCAmelCase				   =       metrics.pop(_A			)
											metrics.update(output.metrics			)
											return PredictionOutput(predictions=predictions.predictions		,			label_ids=predictions.label_ids		,			metrics=_A			)




