

"""simple docstring"""






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


snake_case_ : Tuple   =							logging.get_logger(__name__)

snake_case_ : Optional[int]   =							{
    """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    snake_case__	(			__lowercase  ):







			def __init__(      self      :			List[str]      ,	lowercase      :			Optional[Any]=None      ,	lowercase      :			int=None      ,	*lowercase      :			str      ,	**lowercase      :			List[str]	):





								'''simple docstring'''



								super().__init__(*_lowerCAmelCase      ,	**_lowerCAmelCase	)

								if config is None:
													assert isinstance(self.model      ,	_lowerCAmelCase	), (
													    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
													    f""" {self.model.__class__}"""
													)
													UpperCAmelCase      : int   =					self.model.config
								else:
													UpperCAmelCase      : Optional[Any]   =					config

								UpperCAmelCase      : Optional[Any]   =					data_args
								UpperCAmelCase      : Union[str, Any]   =					self.config.tgt_vocab_size if isinstance(self.config      ,	_lowerCAmelCase	) 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      : List[str]   =					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      : Optional[int]   =					label_smoothed_nll_loss







			def       __lowerCAmelCase					(      self      :			Any      ,	lowercase      :			Dict	):





								'''simple docstring'''



								if self.optimizer is None:
													UpperCAmelCase      : Optional[int]   =					['''bias''', '''LayerNorm.weight''']
													UpperCAmelCase      : Tuple   =					[
													    {
													        '''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      : int   =					Adafactor if self.args.adafactor else AdamW
													if self.args.adafactor:
																		UpperCAmelCase      : Dict   =					Adafactor
																		UpperCAmelCase      : Optional[int]   =					{'''scale_parameter''': False, '''relative_step''': False}
													else:
																		UpperCAmelCase      : Dict   =					AdamW
																		UpperCAmelCase      : Optional[int]   =					{
																		    '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
																		    '''eps''': self.args.adam_epsilon,
																		}
													UpperCAmelCase      : int   =					self.args.learning_rate
													if self.sharded_ddp:
																		UpperCAmelCase      : Union[str, Any]   =					OSS(
																		    params=_lowerCAmelCase      ,	optim=_lowerCAmelCase      ,	**_lowerCAmelCase      ,	)
													else:
																		UpperCAmelCase      : Any   =					optimizer_cls(_lowerCAmelCase      ,	**_lowerCAmelCase	)

								if self.lr_scheduler is None:
													UpperCAmelCase      : Optional[Any]   =					self._get_lr_scheduler(_lowerCAmelCase	)
								else:  # ignoring --lr_scheduler
													logger.warning("scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored."	)







			def       __lowerCAmelCase					(      self      :			Any      ,	lowercase      :			List[Any]	):





								'''simple docstring'''



								UpperCAmelCase      : Any   =					arg_to_scheduler[self.args.lr_scheduler]
								if self.args.lr_scheduler == "constant":
													UpperCAmelCase      : Dict   =					schedule_func(self.optimizer	)
								elif self.args.lr_scheduler == "constant_w_warmup":
													UpperCAmelCase      : Tuple   =					schedule_func(self.optimizer      ,	num_warmup_steps=self.args.warmup_steps	)
								else:
													UpperCAmelCase      : str   =					schedule_func(
													    self.optimizer      ,	num_warmup_steps=self.args.warmup_steps      ,	num_training_steps=_lowerCAmelCase	)
								return scheduler







			def       __lowerCAmelCase					(      self      :			str	):





								'''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       __lowerCAmelCase					(      self      :			Dict      ,	lowercase      :			int      ,	lowercase      :			Union[str, Any]      ,	lowercase      :			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      : Any   =					model(**_lowerCAmelCase      ,	use_cache=_lowerCAmelCase	)[0]
																		UpperCAmelCase      : Union[str, Any]   =					self.loss_fn(logits.view(-1      ,	logits.shape[-1]	)      ,	labels.view(-1	)	)
													else:
																		# compute usual loss via models
																		UpperCAmelCase      : Any   =					model(**_lowerCAmelCase      ,	labels=_lowerCAmelCase      ,	use_cache=_lowerCAmelCase	)[:2]
								else:
													# compute label smoothed loss
													UpperCAmelCase      : str   =					model(**_lowerCAmelCase      ,	use_cache=_lowerCAmelCase	)[0]
													UpperCAmelCase      : Optional[int]   =					torch.nn.functional.log_softmax(_lowerCAmelCase      ,	dim=-1	)
													UpperCAmelCase      : Tuple   =					self.loss_fn(_lowerCAmelCase      ,	_lowerCAmelCase      ,	self.args.label_smoothing      ,	ignore_index=self.config.pad_token_id	)
								return loss, logits







			def       __lowerCAmelCase					(      self      :			Any      ,	lowercase      :			Optional[Any]      ,	lowercase      :			Dict	):





								'''simple docstring'''



								UpperCAmelCase      : Union[str, Any]   =					inputs.pop("labels"	)
								UpperCAmelCase      : Optional[Any]   =					self._compute_loss(_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase	)
								return loss







			def       __lowerCAmelCase					(      self      :			int      ,	lowercase      :			Dict      ,	lowercase      :			Any      ,	lowercase      :			List[Any]      ,	lowercase      :			List[str] = None      ,	):





								'''simple docstring'''



								UpperCAmelCase      : str   =					self._prepare_inputs(_lowerCAmelCase	)

								UpperCAmelCase      : Any   =					{
								    '''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      : Optional[int]   =					self.model.generate(
													    inputs["input_ids"]      ,	attention_mask=inputs["attention_mask"]      ,	**_lowerCAmelCase      ,	)
													# in case the batch is shorter than max length, the output should be padded
													if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																		UpperCAmelCase      : List[Any]   =					self._pad_tensors_to_max_len(_lowerCAmelCase      ,	gen_kwargs["max_length"]	)

								UpperCAmelCase      : Union[str, Any]   =					inputs.pop("labels"	)
								with torch.no_grad():
													# compute loss on predict data
													UpperCAmelCase      : List[Any]   =					self._compute_loss(_lowerCAmelCase      ,	_lowerCAmelCase      ,	_lowerCAmelCase	)

								UpperCAmelCase      : Optional[Any]   =					loss.mean().detach()
								if self.args.prediction_loss_only:
													return (loss, None, None)

								UpperCAmelCase      : int   =					generated_tokens if self.args.predict_with_generate else logits

								if labels.shape[-1] < gen_kwargs["max_length"]:
													UpperCAmelCase      : List[Any]   =					self._pad_tensors_to_max_len(_lowerCAmelCase      ,	gen_kwargs["max_length"]	)

								return (loss, logits, labels)







			def       __lowerCAmelCase					(      self      :			Union[str, Any]      ,	lowercase      :			Tuple      ,	lowercase      :			Any	):





								'''simple docstring'''



								UpperCAmelCase      : str   =					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      : str   =					pad_token_id * torch.ones(
								    (tensor.shape[0], max_length)      ,	dtype=tensor.dtype      ,	device=tensor.device	)
								UpperCAmelCase      : Optional[Any]   =					tensor
								return padded_tensor









import json
import sys
import tempfile
import unittest
from pathlib import Path

import transformers
from transformers import (
    CONFIG_MAPPING,
    IMAGE_PROCESSOR_MAPPING,
    AutoConfig,
    AutoImageProcessor,
    CLIPConfig,
    CLIPImageProcessor,
)
from transformers.testing_utils import DUMMY_UNKNOWN_IDENTIFIER


sys.path.append(str(Path(__file__).parent.parent.parent.parent / """utils"""))

from test_module.custom_configuration import CustomConfig  # noqa E402
from test_module.custom_image_processing import CustomImageProcessor  # noqa E402





class 				a			(      unittest.TestCase			):






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				Optional[int]           =   0






				def 			snake_case_    (  self					):

										"""simple docstring"""




										__SCREAMING_SNAKE_CASE:				List[str]           =   AutoImageProcessor.from_pretrained('''openai/clip-vit-base-patch32'''					)
										self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				List[str]           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				int           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				Tuple           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																__SCREAMING_SNAKE_CASE:				str           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				int           =   CLIPConfig()

																# Create a dummy config file with image_proceesor_type
																__SCREAMING_SNAKE_CASE:				Any           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   Path(_lowerCAmelCase					) / '''config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																# remove image_processor_type to make sure config.json alone is enough to load image processor locally
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					).to_dict()

																config_dict.pop('''image_processor_type'''					)
																__SCREAMING_SNAKE_CASE:				List[str]           =   CLIPImageProcessor(**_lowerCAmelCase					)

																# save in new folder
																model_config.save_pretrained(_lowerCAmelCase					)
																config.save_pretrained(_lowerCAmelCase					)

																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)

																# make sure private variable is not incorrectly saved
																__SCREAMING_SNAKE_CASE:				int           =   json.loads(config.to_json_string()					)
																self.assertTrue('''_processor_class''' not in dict_as_saved					)

										self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with tempfile.TemporaryDirectory() as tmpdirname:
																__SCREAMING_SNAKE_CASE:				Optional[int]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																json.dump(
																    {'''image_processor_type''': '''CLIPImageProcessor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )

																__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, '''clip-base is not a local folder and is not a valid model identifier'''					):
																__SCREAMING_SNAKE_CASE:				Tuple           =   AutoImageProcessor.from_pretrained('''clip-base'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, R'''aaaaaa is not a valid git identifier \(branch name, tag name or commit id\)'''					):
																__SCREAMING_SNAKE_CASE:				Dict           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase			, revision='''aaaaaa'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaisesRegex(
										    _lowerCAmelCase			, '''hf-internal-testing/config-no-model does not appear to have a file named preprocessor_config.json.'''			, ):
																__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/config-no-model'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										with self.assertRaises(_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Optional[Any]           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''					)
										# If remote code is disabled, we can't load this config.
										with self.assertRaises(_lowerCAmelCase					):
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)

										__SCREAMING_SNAKE_CASE:				Optional[Any]           =   AutoImageProcessor.from_pretrained(
										    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
										self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)

										# Test image processor can be reloaded.
										with tempfile.TemporaryDirectory() as tmp_dir:
																image_processor.save_pretrained(_lowerCAmelCase					)
																__SCREAMING_SNAKE_CASE:				List[str]           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase			, trust_remote_code=_lowerCAmelCase					)
										self.assertEqual(reloaded_image_processor.__class__.__name__			, '''NewImageProcessor'''					)






				def 			snake_case_    (  self					):

										"""simple docstring"""




										try:
																AutoConfig.register('''custom'''			, _lowerCAmelCase					)
																AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)
																# Trying to register something existing in the Transformers library will raise an error
																with self.assertRaises(_lowerCAmelCase					):
																						AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)

																with tempfile.TemporaryDirectory() as tmpdirname:
																						__SCREAMING_SNAKE_CASE:				Optional[int]           =   Path(_lowerCAmelCase					) / '''preprocessor_config.json'''
																						__SCREAMING_SNAKE_CASE:				Union[str, Any]           =   Path(_lowerCAmelCase					) / '''config.json'''
																						json.dump(
																						    {'''feature_extractor_type''': '''CLIPFeatureExtractor''', '''processor_class''': '''CLIPProcessor'''}			, open(_lowerCAmelCase			, '''w'''					)			, )
																						json.dump({'''model_type''': '''clip'''}			, open(_lowerCAmelCase			, '''w'''					)					)

																						__SCREAMING_SNAKE_CASE:				List[str]           =   CustomImageProcessor.from_pretrained(_lowerCAmelCase					)

																# Now that the config is registered, it can be used as any other config with the auto-API
																with tempfile.TemporaryDirectory() as tmp_dir:
																						image_processor.save_pretrained(_lowerCAmelCase					)
																						__SCREAMING_SNAKE_CASE:				int           =   AutoImageProcessor.from_pretrained(_lowerCAmelCase					)
																						self.assertIsInstance(_lowerCAmelCase			, _lowerCAmelCase					)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																						del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]






				def 			snake_case_    (  self					):

										"""simple docstring"""




										class 				a			(      __lowercase			):
														SCREAMING_SNAKE_CASE__     :   Tuple		      =     True

										try:
																AutoConfig.register('''custom'''			, _lowerCAmelCase					)
																AutoImageProcessor.register(_lowerCAmelCase			, _lowerCAmelCase					)
																# If remote code is not set, the default is to use local
																__SCREAMING_SNAKE_CASE:				str           =   AutoImageProcessor.from_pretrained('''hf-internal-testing/test_dynamic_image_processor'''					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(image_processor.is_local					)

																# If remote code is disabled, we load the local one.
																__SCREAMING_SNAKE_CASE:				Tuple           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(image_processor.is_local					)

																# If remote is enabled, we load from the Hub
																__SCREAMING_SNAKE_CASE:				Any           =   AutoImageProcessor.from_pretrained(
																    '''hf-internal-testing/test_dynamic_image_processor'''			, trust_remote_code=_lowerCAmelCase					)
																self.assertEqual(image_processor.__class__.__name__			, '''NewImageProcessor'''					)
																self.assertTrue(not hasattr(_lowerCAmelCase			, '''is_local'''					)					)

										finally:
																if "custom" in CONFIG_MAPPING._extra_content:
																						del CONFIG_MAPPING._extra_content["custom"]
																if CustomConfig in IMAGE_PROCESSOR_MAPPING._extra_content:
																						del IMAGE_PROCESSOR_MAPPING._extra_content[CustomConfig]



