






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


__snake_case    :    List[str]         =     get_logger(__name__)







def 	_UpperCamelCase   (    UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					List[str]    ,       UpperCamelCase_       :					List[str]    ,       UpperCamelCase_       :					Tuple=0		)  ->   Any:

						"""simple docstring"""



						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		):
												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(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:
																		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(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:
																		lowerCAmelCase__    						=						os.path.join(UpperCamelCase_    ,       F"{MODEL_NAME}_{model_index}"		)
																		os.makedirs(UpperCamelCase_    ,       exist_ok=UpperCamelCase_		)
																		logger.info(F"Saving model to {ckpt_dir}"		)
																		lowerCAmelCase__    						=						{'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 	_UpperCamelCase   (    UpperCamelCase_       :					List[Any]    ,       UpperCamelCase_       :					List[Any]    ,       UpperCamelCase_       :					List[Any]    ,       UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					str=0		)  ->   List[Any]:

						"""simple docstring"""



						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
																		lowerCAmelCase__    						=						F"{MODEL_NAME}.bin" if model_index == 0 else F"{MODEL_NAME}_{model_index}.bin"
																		lowerCAmelCase__    						=						os.path.join(UpperCamelCase_    ,       UpperCamelCase_		)
																		logger.info(F"Loading model from {input_model_file}"		)
																		lowerCAmelCase__    						=						torch.load(UpperCamelCase_		)
																		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(UpperCamelCase_    ,       UpperCamelCase_		)
																		logger.info(F"Loading model from {input_model_file}"		)
																		lowerCAmelCase__    						=						torch.load(UpperCamelCase_		)
																		logger.info(F"Model loaded from {input_model_file}"		)
												elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
																		lowerCAmelCase__    						=						(
																		    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}"		)
																		lowerCAmelCase__    						=						{'model': model.state_dict()}
																		dist_cp.load_state_dict(
																		    state_dict=UpperCamelCase_    ,       storage_reader=dist_cp.FileSystemReader(UpperCamelCase_		)    ,       planner=DefaultLoadPlanner()    ,       )
																		lowerCAmelCase__    						=						state_dict['model']
																		logger.info(F"Model loaded from {ckpt_dir}"		)
												model.load_state_dict(UpperCamelCase_		)







def 	_UpperCamelCase   (    UpperCamelCase_       :					int    ,       UpperCamelCase_       :					Optional[int]    ,       UpperCamelCase_       :					Any    ,       UpperCamelCase_       :					List[str]    ,       UpperCamelCase_       :					Optional[Any]    ,       UpperCamelCase_       :					Optional[Any]=0		)  ->   Any:

						"""simple docstring"""



						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		):
												lowerCAmelCase__    						=						FSDP.optim_state_dict(UpperCamelCase_    ,       UpperCamelCase_		)
												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(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:
																		lowerCAmelCase__    						=						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 	_UpperCamelCase   (    UpperCamelCase_       :					Dict    ,       UpperCamelCase_       :					str    ,       UpperCamelCase_       :					Tuple    ,       UpperCamelCase_       :					Tuple    ,       UpperCamelCase_       :					int    ,       UpperCamelCase_       :					Union[str, Any]=0		)  ->   Optional[Any]:

						"""simple docstring"""



						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:
																		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(UpperCamelCase_    ,       UpperCamelCase_		)
																		logger.info(F"Loading Optimizer state from {input_optimizer_file}"		)
																		lowerCAmelCase__    						=						torch.load(UpperCamelCase_		)
																		logger.info(F"Optimizer state loaded from {input_optimizer_file}"		)
												else:
																		lowerCAmelCase__    						=						(
																		    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}"		)
																		lowerCAmelCase__    						=						load_sharded_optimizer_state_dict(
																		    model_state_dict=model.state_dict()    ,       optimizer_key='optimizer'    ,       storage_reader=dist_cp.FileSystemReader(UpperCamelCase_		)    ,       )
																		lowerCAmelCase__    						=						optim_state['optimizer']
																		logger.info(F"Optimizer loaded from {ckpt_dir}"		)
												lowerCAmelCase__    						=						FSDP.optim_state_dict_to_load(UpperCamelCase_    ,       UpperCamelCase_    ,       UpperCamelCase_		)
												optimizer.load_state_dict(UpperCamelCase_		)













from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_flax_available, is_torch_available


__snake_case    :    Dict         =     {"""configuration_speech_encoder_decoder""": ["""SpeechEncoderDecoderConfig"""]}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__snake_case    :    Union[str, Any]         =     ["""SpeechEncoderDecoderModel"""]

try:
					if not is_flax_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					__snake_case    :    str         =     ["""FlaxSpeechEncoderDecoderModel"""]

if TYPE_CHECKING:
					from .configuration_speech_encoder_decoder import SpeechEncoderDecoderConfig

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_speech_encoder_decoder import SpeechEncoderDecoderModel

					try:
										if not is_flax_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_flax_speech_encoder_decoder import FlaxSpeechEncoderDecoderModel

else:
					import sys

					__snake_case    :    Tuple         =     _LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)





