



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


lowercase_													=      logging.get_logger(__name__)

lowercase_													=      {
    '''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			(  SCREAMING_SNAKE_CASE					):



	def __init__(    self				:     Optional[int]					,						snake_case_				:     int=None					,						snake_case_				:     List[Any]=None					,						*snake_case_				:     str					,						**snake_case_				:     str)->				Union[str, Any]:
						super().__init__(*snake_case_					,						**snake_case_)

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

						__lowerCAmelCase						 =data_args
						__lowerCAmelCase						 =self.config.tgt_vocab_size if isinstance(self.config					,						snake_case_) 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:
											__lowerCAmelCase						 =torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)
						else:
											# dynamically import label_smoothed_nll_loss
											from utils import label_smoothed_nll_loss

											__lowerCAmelCase						 =label_smoothed_nll_loss



	def 		UpperCamelCase				(    self				:     int					,						snake_case_				:     int)->				Optional[int]:
						if self.optimizer is None:
											__lowerCAmelCase						 =["""bias""", """LayerNorm.weight"""]
											__lowerCAmelCase						 =[
											    {
											        """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,
											    },
											]
											__lowerCAmelCase						 =Adafactor if self.args.adafactor else AdamW
											if self.args.adafactor:
																__lowerCAmelCase						 =Adafactor
																__lowerCAmelCase						 ={"""scale_parameter""": False, """relative_step""": False}
											else:
																__lowerCAmelCase						 =AdamW
																__lowerCAmelCase						 ={
																    """betas""": (self.args.adam_betaa, self.args.adam_betaa),
																    """eps""": self.args.adam_epsilon,
																}
											__lowerCAmelCase						 =self.args.learning_rate
											if self.sharded_ddp:
																__lowerCAmelCase						 =OSS(
																    params=snake_case_					,						optim=snake_case_					,						**snake_case_					,						)
											else:
																__lowerCAmelCase						 =optimizer_cls(snake_case_					,						**snake_case_)

						if self.lr_scheduler is None:
											__lowerCAmelCase						 =self._get_lr_scheduler(snake_case_)
						else:  # ignoring --lr_scheduler
											logger.warning("""scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.""")



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     Union[str, Any])->				str:
						__lowerCAmelCase						 =arg_to_scheduler[self.args.lr_scheduler]
						if self.args.lr_scheduler == "constant":
											__lowerCAmelCase						 =schedule_func(self.optimizer)
						elif self.args.lr_scheduler == "constant_w_warmup":
											__lowerCAmelCase						 =schedule_func(self.optimizer					,						num_warmup_steps=self.args.warmup_steps)
						else:
											__lowerCAmelCase						 =schedule_func(
											    self.optimizer					,						num_warmup_steps=self.args.warmup_steps					,						num_training_steps=snake_case_)
						return scheduler



	def 		UpperCamelCase				(    self				:     str)->				Optional[torch.utils.data.Sampler]:
						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				:     str					,						snake_case_				:     int					,						snake_case_				:     Optional[int]					,						snake_case_				:     List[str])->				Any:
						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
																__lowerCAmelCase						 =model(**snake_case_					,						use_cache=snake_case_)[0]
																__lowerCAmelCase						 =self.loss_fn(logits.view(-1					,						logits.shape[-1])					,						labels.view(-1))
											else:
																# compute usual loss via models
																__lowerCAmelCase ,						__lowerCAmelCase						 =model(**snake_case_					,						labels=snake_case_					,						use_cache=snake_case_)[:2]
						else:
											# compute label smoothed loss
											__lowerCAmelCase						 =model(**snake_case_					,						use_cache=snake_case_)[0]
											__lowerCAmelCase						 =torch.nn.functional.log_softmax(snake_case_					,						dim=-1)
											__lowerCAmelCase ,						__lowerCAmelCase						 =self.loss_fn(snake_case_					,						snake_case_					,						self.args.label_smoothing					,						ignore_index=self.config.pad_token_id)
						return loss, logits



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     int					,						snake_case_				:     List[str])->				Union[str, Any]:
						__lowerCAmelCase						 =inputs.pop("""labels""")
						__lowerCAmelCase ,						__lowerCAmelCase						 =self._compute_loss(snake_case_					,						snake_case_					,						snake_case_)
						return loss



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     nn.Module					,						snake_case_				:     Dict[str, Union[torch.Tensor, Any]]					,						snake_case_				:     bool					,						snake_case_				:     Optional[List[str]] = None					,						)->				Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
						__lowerCAmelCase						 =self._prepare_inputs(snake_case_)

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

						__lowerCAmelCase						 =inputs.pop("""labels""")
						with torch.no_grad():
											# compute loss on predict data
											__lowerCAmelCase ,						__lowerCAmelCase						 =self._compute_loss(snake_case_					,						snake_case_					,						snake_case_)

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

						__lowerCAmelCase						 =generated_tokens if self.args.predict_with_generate else logits

						if labels.shape[-1] < gen_kwargs["max_length"]:
											__lowerCAmelCase						 =self._pad_tensors_to_max_len(snake_case_					,						gen_kwargs["""max_length"""])

						return (loss, logits, labels)



	def 		UpperCamelCase				(    self				:     Optional[Any]					,						snake_case_				:     Union[str, Any]					,						snake_case_				:     Optional[int])->				Any:
						# If PAD token is not defined at least EOS token has to be defined
						__lowerCAmelCase						 =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}""")

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









from __future__ import annotations

import unittest

import numpy as np

from transformers import BlipTextConfig
from transformers.testing_utils import require_tf, slow
from transformers.utils import is_tf_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask


if is_tf_available():
				import tensorflow as tf

				from transformers import TFBlipTextModel
				from transformers.models.blip.modeling_tf_blip import TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST






class     __a			:



	def __init__(    self				:     int					,						snake_case_				:     Dict					,						snake_case_				:     List[str]=12					,						snake_case_				:     Optional[int]=7					,						snake_case_				:     Optional[Any]=True					,						snake_case_				:     int=True					,						snake_case_				:     int=True					,						snake_case_				:     Any=99					,						snake_case_				:     List[Any]=32					,						snake_case_				:     List[Any]=32					,						snake_case_				:     Any=2					,						snake_case_				:     Optional[Any]=4					,						snake_case_				:     Optional[Any]=37					,						snake_case_				:     Tuple=0.1					,						snake_case_				:     str=0.1					,						snake_case_				:     Any=5_12					,						snake_case_				:     Optional[Any]=0.0_2					,						snake_case_				:     List[Any]=0					,						snake_case_				:     Union[str, Any]=None					,						)->				List[Any]:
						__lowerCAmelCase						 =parent
						__lowerCAmelCase						 =batch_size
						__lowerCAmelCase						 =seq_length
						__lowerCAmelCase						 =is_training
						__lowerCAmelCase						 =use_input_mask
						__lowerCAmelCase						 =use_labels
						__lowerCAmelCase						 =vocab_size
						__lowerCAmelCase						 =hidden_size
						__lowerCAmelCase						 =projection_dim
						__lowerCAmelCase						 =num_hidden_layers
						__lowerCAmelCase						 =num_attention_heads
						__lowerCAmelCase						 =intermediate_size
						__lowerCAmelCase						 =dropout
						__lowerCAmelCase						 =attention_dropout
						__lowerCAmelCase						 =max_position_embeddings
						__lowerCAmelCase						 =initializer_range
						__lowerCAmelCase						 =scope
						__lowerCAmelCase						 =bos_token_id



	def 		UpperCamelCase				(    self				:     List[Any])->				Optional[int]:
						__lowerCAmelCase						 =ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size)

						__lowerCAmelCase						 =None
						if self.use_input_mask:
											__lowerCAmelCase						 =random_attention_mask([self.batch_size, self.seq_length])

						if input_mask is not None:
											__lowerCAmelCase						 =input_mask.numpy()
											__lowerCAmelCase ,						__lowerCAmelCase						 =input_mask.shape
											__lowerCAmelCase						 =np.random.randint(1					,						seq_length - 1					,						size=(batch_size,))
											for batch_idx, start_index in enumerate(snake_case_):
																__lowerCAmelCase						 =1
																__lowerCAmelCase						 =0

						__lowerCAmelCase						 =self.get_config()

						return config, input_ids, tf.convert_to_tensor(snake_case_)



	def 		UpperCamelCase				(    self				:     List[str])->				str:
						return BlipTextConfig(
						    vocab_size=self.vocab_size					,						hidden_size=self.hidden_size					,						projection_dim=self.projection_dim					,						num_hidden_layers=self.num_hidden_layers					,						num_attention_heads=self.num_attention_heads					,						intermediate_size=self.intermediate_size					,						dropout=self.dropout					,						attention_dropout=self.attention_dropout					,						max_position_embeddings=self.max_position_embeddings					,						initializer_range=self.initializer_range					,						bos_token_id=self.bos_token_id					,						)



	def 		UpperCamelCase				(    self				:     Union[str, Any]					,						snake_case_				:     Tuple					,						snake_case_				:     List[Any]					,						snake_case_				:     List[Any])->				List[str]:
						__lowerCAmelCase						 =TFBlipTextModel(config=snake_case_)
						__lowerCAmelCase						 =model(snake_case_					,						attention_mask=snake_case_					,						training=snake_case_)
						__lowerCAmelCase						 =model(snake_case_					,						training=snake_case_)
						self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size))
						self.parent.assertEqual(result.pooler_output.shape					,						(self.batch_size, self.hidden_size))



	def 		UpperCamelCase				(    self				:     Dict)->				Tuple:
						__lowerCAmelCase						 =self.prepare_config_and_inputs()
						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase						 =config_and_inputs
						__lowerCAmelCase						 ={"""input_ids""": input_ids, """attention_mask""": input_mask}
						return config, inputs_dict






@require_tf
class     __a			(  SCREAMING_SNAKE_CASE						,     unittest.TestCase					):
	SCREAMING_SNAKE_CASE  	=							(TFBlipTextModel,) if is_tf_available() else ()
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False
	SCREAMING_SNAKE_CASE  	=							False



	def 		UpperCamelCase				(    self				:     Optional[Any])->				Any:
						__lowerCAmelCase						 =BlipTextModelTester(self)
						__lowerCAmelCase						 =ConfigTester(self					,						config_class=snake_case_					,						hidden_size=37)



	def 		UpperCamelCase				(    self				:     List[Any])->				List[str]:
						self.config_tester.run_common_tests()



	def 		UpperCamelCase				(    self				:     List[str])->				Dict:
						__lowerCAmelCase						 =self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*snake_case_)



	def 		UpperCamelCase				(    self				:     Optional[Any])->				str:
						pass



	def 		UpperCamelCase				(    self				:     List[str])->				int:
						pass



	@unittest.skip(reason="""Blip does not use inputs_embeds""")
	def 		UpperCamelCase				(    self				:     Optional[Any])->				int:
						pass



	@unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING""")
	def 		UpperCamelCase				(    self				:     List[Any])->				str:
						pass



	@unittest.skip(reason="""BlipTextModel has no base class and is not available in MODEL_MAPPING""")
	def 		UpperCamelCase				(    self				:     Optional[int])->				Tuple:
						pass



	@slow
	def 		UpperCamelCase				(    self				:     int)->				str:
						for model_name in TF_BLIP_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											__lowerCAmelCase						 =TFBlipTextModel.from_pretrained(snake_case_)
											self.assertIsNotNone(snake_case_)



	def 		UpperCamelCase				(    self				:     str					,						snake_case_				:     List[str]=True)->				int:
						super().test_pt_tf_model_equivalence(allow_missing_keys=snake_case_)




