


from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
		from fairscale.optim import OSS


lowerCAmelCase__    :Any       =				logging.get_logger(__name__)

lowerCAmelCase__    :Any       =				{
    """linear""": get_linear_schedule_with_warmup,
    """cosine""": get_cosine_schedule_with_warmup,
    """cosine_w_restarts""": get_cosine_with_hard_restarts_schedule_with_warmup,
    """polynomial""": get_polynomial_decay_schedule_with_warmup,
    """constant""": get_constant_schedule,
    """constant_w_warmup""": get_constant_schedule_with_warmup,
}

class 						__a (							a_     ):
		def __init__(     self     ,       _SCREAMING_SNAKE_CASE=None     ,       _SCREAMING_SNAKE_CASE=None     ,       *_SCREAMING_SNAKE_CASE     ,       **_SCREAMING_SNAKE_CASE		)					->			List[str]:




									"""simple docstring"""





									super().__init__(*_lowercase     ,       **_lowercase		)

									if config is None:
																assert isinstance(self.model     ,       _lowercase		), (
																    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
																    f''' {self.model.__class__}'''
																)
																_UpperCAmelCase											=     self.model.config
									else:
																_UpperCAmelCase											=     config

									_UpperCAmelCase											=     data_args
									_UpperCAmelCase											=     self.config.tgt_vocab_size if isinstance(self.config     ,       _lowercase		) else self.config.vocab_size

									if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
																assert self.config.pad_token_id is not None, (
																    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
																    " calculation or doing label smoothing."
																)

									if self.config.pad_token_id is None and self.config.eos_token_id is not None:
																logger.warning(
																    f'''The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for'''
																    ' padding..'		)

									if self.args.label_smoothing == 0:
																_UpperCAmelCase											=     torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id		)
									else:
																# dynamically import label_smoothed_nll_loss
																from utils import label_smoothed_nll_loss

																_UpperCAmelCase											=     label_smoothed_nll_loss
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			Dict:




									"""simple docstring"""





									if self.optimizer is None:
																_UpperCAmelCase											=     ["""bias""", """LayerNorm.weight"""]
																_UpperCAmelCase											=     [
																    {
																        """params""": [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay		)],
																        """weight_decay""": self.args.weight_decay,
																    },
																    {
																        """params""": [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay		)],
																        """weight_decay""": 0.0,
																    },
																]
																_UpperCAmelCase											=     Adafactor if self.args.adafactor else AdamW
																if self.args.adafactor:
																							_UpperCAmelCase											=     Adafactor
																							_UpperCAmelCase											=     {"""scale_parameter""": False, """relative_step""": False}
																else:
																							_UpperCAmelCase											=     AdamW
																							_UpperCAmelCase											=     {
																							    """betas""": (self.args.adam_betaa, self.args.adam_betaa),
																							    """eps""": self.args.adam_epsilon,
																							}
																_UpperCAmelCase											=     self.args.learning_rate
																if self.sharded_ddp:
																							_UpperCAmelCase											=     OSS(
																							    params=_lowercase     ,       optim=_lowercase     ,       **_lowercase     ,       )
																else:
																							_UpperCAmelCase											=     optimizer_cls(_lowercase     ,       **_lowercase		)

									if self.lr_scheduler is None:
																_UpperCAmelCase											=     self._get_lr_scheduler(_lowercase		)
									else:  # ignoring --lr_scheduler
																logger.warning('scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'		)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE		)					->			Dict:




									"""simple docstring"""





									_UpperCAmelCase											=     arg_to_scheduler[self.args.lr_scheduler]
									if self.args.lr_scheduler == "constant":
																_UpperCAmelCase											=     schedule_func(self.optimizer		)
									elif self.args.lr_scheduler == "constant_w_warmup":
																_UpperCAmelCase											=     schedule_func(self.optimizer     ,       num_warmup_steps=self.args.warmup_steps		)
									else:
																_UpperCAmelCase											=     schedule_func(
																    self.optimizer     ,       num_warmup_steps=self.args.warmup_steps     ,       num_training_steps=_lowercase		)
									return scheduler
		def       UpperCAmelCase__							(     self		)					->			Optional[torch.utils.data.Sampler]:




									"""simple docstring"""





									if isinstance(self.train_dataset     ,       torch.utils.data.IterableDataset		):
																return None
									elif is_torch_tpu_available():
																return get_tpu_sampler(self.train_dataset		)
									else:
																if self.args.sortish_sampler:
																							self.train_dataset.make_sortish_sampler(
																							    self.args.per_device_train_batch_size     ,       distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)     ,       )

																return (
																    RandomSampler(self.train_dataset		)
																    if self.args.local_rank == -1
																    else DistributedSampler(self.train_dataset		)
																)
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Dict:




									"""simple docstring"""





									if self.args.label_smoothing == 0:
																if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
																							# force training to ignore pad token
																							_UpperCAmelCase											=     model(**_lowercase     ,       use_cache=_lowercase		)[0]
																							_UpperCAmelCase											=     self.loss_fn(logits.view(-1     ,       logits.shape[-1]		)     ,       labels.view(-1		)		)
																else:
																							# compute usual loss via models
																							_UpperCAmelCase											=     model(**_lowercase     ,       labels=_lowercase     ,       use_cache=_lowercase		)[:2]
									else:
																# compute label smoothed loss
																_UpperCAmelCase											=     model(**_lowercase     ,       use_cache=_lowercase		)[0]
																_UpperCAmelCase											=     torch.nn.functional.log_softmax(_lowercase     ,       dim=-1		)
																_UpperCAmelCase											=     self.loss_fn(_lowercase     ,       _lowercase     ,       self.args.label_smoothing     ,       ignore_index=self.config.pad_token_id		)
									return loss, logits
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Optional[Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     inputs.pop('labels'		)
									_UpperCAmelCase											=     self._compute_loss(_lowercase     ,       _lowercase     ,       _lowercase		)
									return loss
		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE = None     ,       )					->			Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:




									"""simple docstring"""





									_UpperCAmelCase											=     self._prepare_inputs(_lowercase		)

									_UpperCAmelCase											=     {
									    """max_length""": self.data_args.val_max_target_length
									    if self.data_args is not None
									    else self.config.max_length,
									    """num_beams""": self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
									}

									if self.args.predict_with_generate and not self.args.prediction_loss_only:
																_UpperCAmelCase											=     self.model.generate(
																    inputs['input_ids']     ,       attention_mask=inputs['attention_mask']     ,       **_lowercase     ,       )
																# in case the batch is shorter than max length, the output should be padded
																if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																							_UpperCAmelCase											=     self._pad_tensors_to_max_len(_lowercase     ,       gen_kwargs['max_length']		)

									_UpperCAmelCase											=     inputs.pop('labels'		)
									with torch.no_grad():
																# compute loss on predict data
																_UpperCAmelCase											=     self._compute_loss(_lowercase     ,       _lowercase     ,       _lowercase		)

									_UpperCAmelCase											=     loss.mean().detach()
									if self.args.prediction_loss_only:
																return (loss, None, None)

									_UpperCAmelCase											=     generated_tokens if self.args.predict_with_generate else logits

									if labels.shape[-1] < gen_kwargs["max_length"]:
																_UpperCAmelCase											=     self._pad_tensors_to_max_len(_lowercase     ,       gen_kwargs['max_length']		)

									return (loss, logits, labels)






		def       UpperCAmelCase__							(     self     ,       _SCREAMING_SNAKE_CASE     ,       _SCREAMING_SNAKE_CASE		)					->			Union[str, Any]:




									"""simple docstring"""





									_UpperCAmelCase											=     self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

									if pad_token_id is None:
																raise ValueError(
																    'Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'
																    f''' padded to `max_length`={max_length}'''		)

									_UpperCAmelCase											=     pad_token_id * torch.ones(
									    (tensor.shape[0], max_length)     ,       dtype=tensor.dtype     ,       device=tensor.device		)
									_UpperCAmelCase											=     tensor
									return padded_tensor



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        A__:
						"""simple docstring"""






						def __init__(    self  ,      _lowercase  ,      _lowercase=99  ,      _lowercase=13  ,      _lowercase=16  ,      _lowercase=7  ,      _lowercase=True  ,      _lowercase=True  ,      _lowercase=True  ,      _lowercase=False  ,      _lowercase=True  ,      _lowercase=2  ,      _lowercase=32  ,      _lowercase=4  ,      _lowercase=4  ,      _lowercase=30  ,      _lowercase=0  ,      _lowercase=1  ,      _lowercase=2  ,      _lowercase=None  ,      )     -> List[str]:
									a_						:						Union[str, Any] 							= parent
									a_						:						Dict 							= batch_size
									a_						:						Any 							= decoder_seq_length
									# For common tests
									a_						:						Any 							= self.decoder_seq_length
									a_						:						int 							= is_training
									a_						:						List[Any] 							= use_attention_mask
									a_						:						List[str] 							= use_labels

									a_						:						Optional[Any] 							= vocab_size
									a_						:						Dict 							= d_model
									a_						:						Any 							= d_model
									a_						:						List[Any] 							= decoder_layers
									a_						:						List[str] 							= decoder_layers
									a_						:						Union[str, Any] 							= decoder_ffn_dim
									a_						:						Optional[Any] 							= decoder_attention_heads
									a_						:						Optional[Any] 							= decoder_attention_heads
									a_						:						Tuple 							= eos_token_id
									a_						:						Union[str, Any] 							= bos_token_id
									a_						:						Dict 							= pad_token_id
									a_						:						Union[str, Any] 							= decoder_start_token_id
									a_						:						List[Any] 							= use_cache
									a_						:						List[str] 							= max_position_embeddings

									a_						:						int 							= None
									a_						:						Any 							= decoder_seq_length
									a_						:						Union[str, Any] 							= 2
									a_						:						Optional[int] 							= 1





						def 			UpperCamelCase__					(    self    )     -> int:
									a_						:						Optional[int] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      self.vocab_size    )

									a_						:						Optional[int] 							= None
									if self.use_attention_mask:
												a_						:						List[Any] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      vocab_size=2    )

									a_						:						List[Any] 							= None
									if self.use_labels:
												a_						:						Union[str, Any] 							= ids_tensor([self.batch_size, self.decoder_seq_length]  ,      self.vocab_size    )

									a_						:						Any 							= 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 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      _lowercase  ,      )     -> List[Any]:
									a_						:						Optional[int] 							= True
									a_						:						List[str] 							= TrOCRDecoder(config=_lowercase    ).to(_lowercase    ).eval()
									a_						:						str 							= input_ids[:2]

									input_ids[input_ids == 0] += 1
									# first forward pass
									a_						:						Dict 							= model(_lowercase  ,      use_cache=_lowercase    )
									a_						:						Union[str, Any] 							= model(_lowercase    )
									a_						:						Union[str, Any] 							= model(_lowercase  ,      use_cache=_lowercase    )

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

									a_						:						Dict 							= outputs["""past_key_values"""]

									# create hypothetical next token and extent to next_input_ids
									a_						:						List[str] 							= ids_tensor((2, 1)  ,      config.vocab_size - 1    ) + 1

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

									a_						:						Dict 							= model(_lowercase    )["""last_hidden_state"""]
									a_						:						List[str] 							= model(_lowercase  ,      past_key_values=_lowercase    )["""last_hidden_state"""]

									# select random slice
									a_						:						Optional[Any] 							= ids_tensor((1,)  ,      output_from_past.shape[-1]    ).item()
									a_						:						Union[str, Any] 							= output_from_no_past[:, next_input_ids.shape[-1] - 1, random_slice_idx].detach()
									a_						:						Any 							= output_from_past[:, 0, random_slice_idx].detach()

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





						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									a_						:						List[Any] 							= self.prepare_config_and_inputs()
									a_							,					a_							,					a_							,					a_						:						List[Any] 							= config_and_inputs

									a_						:						str 							= {"""input_ids""": input_ids, """attention_mask""": attention_mask}
									return config, inputs_dict







@require_torch
class        A__(a_,  a_,  a_,  unittest.TestCase							):
						"""simple docstring"""






						_A    :				Optional[Any]       					=			(TrOCRDecoder, TrOCRForCausalLM) if is_torch_available() else ()
						_A    :				Optional[int]       					=			(TrOCRForCausalLM,) if is_torch_available() else ()
						_A    :				Tuple       					=			{'''text-generation''': TrOCRForCausalLM} if is_torch_available() else {}
						_A    :				List[str]       					=			True
						_A    :				List[str]       					=			False





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						str 							= TrOCRStandaloneDecoderModelTester(self  ,      is_training=_lowercase    )
									a_						:						str 							= ConfigTester(self  ,      config_class=_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									pass





						def 			UpperCamelCase__					(    self    )     -> int:
									pass





						def 			UpperCamelCase__					(    self    )     -> str:
									pass





						def 			UpperCamelCase__					(    self    )     -> Dict:
									self.config_tester.run_common_tests()





						def 			UpperCamelCase__					(    self    )     -> List[str]:
									a_						:						int 							= self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_decoder_model_past(*_lowercase    )





						def 			UpperCamelCase__					(    self    )     -> Tuple:
									return





						@unittest.skip("""The model doesn't support left padding"""    )  # and it's not used enough to be worth fixing :)
						def 			UpperCamelCase__					(    self    )     -> Optional[Any]:
									pass


