
"""simple docstring"""




import unittest

from transformers import TrOCRConfig
from transformers.testing_utils import is_torch_available, require_torch, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers.models.trocr.modeling_trocr import TrOCRDecoder, TrOCRForCausalLM





@require_torch
class 			snake_case  :

							def __init__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=99							,							SCREAMING_SNAKE_CASE_=13							,							SCREAMING_SNAKE_CASE_=16							,							SCREAMING_SNAKE_CASE_=7							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=2							,							SCREAMING_SNAKE_CASE_=32							,							SCREAMING_SNAKE_CASE_=4							,							SCREAMING_SNAKE_CASE_=4							,							SCREAMING_SNAKE_CASE_=30							,							SCREAMING_SNAKE_CASE_=0							,							SCREAMING_SNAKE_CASE_=1							,							SCREAMING_SNAKE_CASE_=2							,							SCREAMING_SNAKE_CASE_=None							,							):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			parent
									SCREAMING_SNAKE_CASE_    			=			batch_size
									SCREAMING_SNAKE_CASE_    			=			decoder_seq_length
									# For common tests
									SCREAMING_SNAKE_CASE_    			=			self.decoder_seq_length
									SCREAMING_SNAKE_CASE_    			=			is_training
									SCREAMING_SNAKE_CASE_    			=			use_attention_mask
									SCREAMING_SNAKE_CASE_    			=			use_labels

									SCREAMING_SNAKE_CASE_    			=			vocab_size
									SCREAMING_SNAKE_CASE_    			=			d_model
									SCREAMING_SNAKE_CASE_    			=			d_model
									SCREAMING_SNAKE_CASE_    			=			decoder_layers
									SCREAMING_SNAKE_CASE_    			=			decoder_layers
									SCREAMING_SNAKE_CASE_    			=			decoder_ffn_dim
									SCREAMING_SNAKE_CASE_    			=			decoder_attention_heads
									SCREAMING_SNAKE_CASE_    			=			decoder_attention_heads
									SCREAMING_SNAKE_CASE_    			=			eos_token_id
									SCREAMING_SNAKE_CASE_    			=			bos_token_id
									SCREAMING_SNAKE_CASE_    			=			pad_token_id
									SCREAMING_SNAKE_CASE_    			=			decoder_start_token_id
									SCREAMING_SNAKE_CASE_    			=			use_cache
									SCREAMING_SNAKE_CASE_    			=			max_position_embeddings

									SCREAMING_SNAKE_CASE_    			=			None
									SCREAMING_SNAKE_CASE_    			=			decoder_seq_length
									SCREAMING_SNAKE_CASE_    			=			2
									SCREAMING_SNAKE_CASE_    			=			1

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			ids_tensor([self.batch_size, self.decoder_seq_length]							,							self.vocab_size     )

									SCREAMING_SNAKE_CASE_    			=			None
									if self.use_attention_mask:
											SCREAMING_SNAKE_CASE_    			=			ids_tensor([self.batch_size, self.decoder_seq_length]							,							vocab_size=2     )

									SCREAMING_SNAKE_CASE_    			=			None
									if self.use_labels:
											SCREAMING_SNAKE_CASE_    			=			ids_tensor([self.batch_size, self.decoder_seq_length]							,							self.vocab_size     )

									SCREAMING_SNAKE_CASE_    			=			TrOCRConfig(
									    vocab_size=self.vocab_size							,							d_model=self.d_model							,							decoder_layers=self.decoder_layers							,							decoder_ffn_dim=self.decoder_ffn_dim							,							decoder_attention_heads=self.decoder_attention_heads							,							eos_token_id=self.eos_token_id							,							bos_token_id=self.bos_token_id							,							use_cache=self.use_cache							,							pad_token_id=self.pad_token_id							,							decoder_start_token_id=self.decoder_start_token_id							,							max_position_embeddings=self.max_position_embeddings							,							)

									return (config, input_ids, attention_mask, lm_labels)

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			True
									SCREAMING_SNAKE_CASE_    			=			TrOCRDecoder(config=SCREAMING_SNAKE_CASE_     ).to(SCREAMING_SNAKE_CASE_     ).eval()
									SCREAMING_SNAKE_CASE_    			=			input_ids[:2]

									input_ids[input_ids == 0] += 1
									# first forward pass
									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							use_cache=SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							use_cache=SCREAMING_SNAKE_CASE_     )

									self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_     ) == len(SCREAMING_SNAKE_CASE_     )     )
									self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_     ) == len(SCREAMING_SNAKE_CASE_     ) + 1     )

									SCREAMING_SNAKE_CASE_    			=			outputs['''past_key_values''']

									# create hypothetical next token and extent to next_input_ids
									SCREAMING_SNAKE_CASE_    			=			ids_tensor((2, 1)							,							config.vocab_size - 1     ) + 1

									# append to next input_ids and
									SCREAMING_SNAKE_CASE_    			=			torch.cat([input_ids, next_tokens]							,							dim=-1     )

									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_     )['''last_hidden_state''']
									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							past_key_values=SCREAMING_SNAKE_CASE_     )['''last_hidden_state''']

									# select random slice
									SCREAMING_SNAKE_CASE_    			=			ids_tensor((1,)							,							output_from_past.shape[-1]     ).item()
									SCREAMING_SNAKE_CASE_    			=			output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
									SCREAMING_SNAKE_CASE_    			=			output_from_past[:, 0, random_slice_idx].detach()

									# test that outputs are equal for slice
									assert torch.allclose(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							atol=1e-3     )







							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.prepare_config_and_inputs()
									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			config_and_inputs

									SCREAMING_SNAKE_CASE_    			=			{'''input_ids''': input_ids, '''attention_mask''': attention_mask}
									return config, inputs_dict







@require_torch
class 			snake_case  (				__lowercase							,      __lowercase							,      __lowercase							,      unittest.TestCase ):
							UpperCAmelCase__      =   (TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
							UpperCAmelCase__      =   (TrOCRForCausalLM,) if is_torch_available() else ()
							UpperCAmelCase__      =   {'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
							UpperCAmelCase__      =   True
							UpperCAmelCase__      =   False

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			TrOCRStandaloneDecoderModelTester(self							,							is_training=SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			ConfigTester(self							,							config_class=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									pass

							def _lowercase       (self     ):

									"""simple docstring"""







									pass

							def _lowercase       (self     ):

									"""simple docstring"""







									pass

							def _lowercase       (self     ):

									"""simple docstring"""







									self.config_tester.run_common_tests()

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_decoder_model_past(*SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									return







							@unittest.skip('''The model doesn\'t support left padding'''     )  # and it's not used enough to be worth fixing :)
							def _lowercase       (self     ):

									"""simple docstring"""







									pass

"""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


lowerCAmelCase__				=					get_logger(__name__)



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			model.state_dict()
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						if accelerator.process_index == 0:
								logger.info(F'Saving model to {output_model_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
						    if model_index == 0
						    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Saving model to {output_model_file}'     )
						torch.save(__a,				__a     )
						logger.info(F'Model saved to {output_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving model to {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': state_dict}

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



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				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(__a     ) != 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
						SCREAMING_SNAKE_CASE_    			=			F'{MODEL_NAME}.bin' if model_index == 0 else F'{MODEL_NAME}_{model_index}.bin'
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.LOCAL_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{MODEL_NAME}_rank{accelerator.process_index}.bin'
						    if model_index == 0
						    else F'{MODEL_NAME}_{model_index}_rank{accelerator.process_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading model from {input_model_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Model loaded from {input_model_file}'     )
				elif fsdp_plugin.state_dict_type == StateDictType.SHARDED_STATE_DICT:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{MODEL_NAME}_{model_index}'     )
						    if F'{MODEL_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading model from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			{'''model''': model.state_dict()}
						dist_cp.load_state_dict(
						    state_dict=__a,				storage_reader=dist_cp.FileSystemReader(__a     ),				planner=DefaultLoadPlanner(),				)
						SCREAMING_SNAKE_CASE_    			=			state_dict['''model''']
						logger.info(F'Model loaded from {ckpt_dir}'     )
				model.load_state_dict(__a     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		os.makedirs(__a,				exist_ok=__a     )
		with FSDP.state_dict_type(
		    __a,				fsdp_plugin.state_dict_type,				fsdp_plugin.state_dict_config,				fsdp_plugin.optim_state_dict_config     ):
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict(__a,				__a     )
				if fsdp_plugin.state_dict_type == StateDictType.FULL_STATE_DICT:
						if accelerator.process_index == 0:
								SCREAMING_SNAKE_CASE_    			=			(
								    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
								)
								SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
								logger.info(F'Saving Optimizer state to {output_optimizer_file}'     )
								torch.save(__a,				__a     )
								logger.info(F'Optimizer state saved in {output_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						os.makedirs(__a,				exist_ok=__a     )
						logger.info(F'Saving Optimizer state to {ckpt_dir}'     )
						dist_cp.save_state_dict(
						    state_dict={'''optimizer''': optim_state},				storage_writer=dist_cp.FileSystemWriter(__a     ),				planner=DefaultSavePlanner(),				)
						logger.info(F'Optimizer state saved in {ckpt_dir}'     )



def 							_lowerCamelCase			(			__a,				__a,				__a,				__a,				__a,				__a=0     ):
		accelerator.wait_for_everyone()
		with FSDP.state_dict_type(
		    __a,				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:
						SCREAMING_SNAKE_CASE_    			=			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:
						SCREAMING_SNAKE_CASE_    			=			(
						    F'{OPTIMIZER_NAME}.bin' if optimizer_index == 0 else F'{OPTIMIZER_NAME}_{optimizer_index}.bin'
						)
						SCREAMING_SNAKE_CASE_    			=			os.path.join(__a,				__a     )
						logger.info(F'Loading Optimizer state from {input_optimizer_file}'     )
						SCREAMING_SNAKE_CASE_    			=			torch.load(__a     )
						logger.info(F'Optimizer state loaded from {input_optimizer_file}'     )
				else:
						SCREAMING_SNAKE_CASE_    			=			(
						    os.path.join(__a,				F'{OPTIMIZER_NAME}_{optimizer_index}'     )
						    if F'{OPTIMIZER_NAME}' not in input_dir
						    else input_dir
						)
						logger.info(F'Loading Optimizer from {ckpt_dir}'     )
						SCREAMING_SNAKE_CASE_    			=			load_sharded_optimizer_state_dict(
						    model_state_dict=model.state_dict(),				optimizer_key='''optimizer''',				storage_reader=dist_cp.FileSystemReader(__a     ),				)
						SCREAMING_SNAKE_CASE_    			=			optim_state['''optimizer''']
						logger.info(F'Optimizer loaded from {ckpt_dir}'     )
				SCREAMING_SNAKE_CASE_    			=			FSDP.optim_state_dict_to_load(__a,				__a,				__a     )
				optimizer.load_state_dict(__a     )