





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    : str    	=  logging.get_logger(__name__)

lowercase    : Optional[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(						__UpperCAmelCase ):

     def __init__(      self,       A=None,       A=None,       *A,       **A      ):





            """simple docstring"""

            super().__init__(*A,       **A      )

            if config is None:
                   assert isinstance(self.model,       A      ), (
                       "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,       A      ) 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,       A      ):





            """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=A,       optim=A,       **A,       )
                   else:
                          _UpperCamelCase              =						optimizer_cls(A,       **A      )

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

     def 	_UpperCamelCase     (      self,       A      ):





            """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=A      )
            return scheduler

     def 	_UpperCamelCase     (      self      ):





            """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,       A,       A,       A      ):





            """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(**A,       use_cache=A      )[0]
                          _UpperCamelCase              =						self.loss_fn(logits.view(-1,       logits.shape[-1]      ),       labels.view(-1      )      )
                   else:
                          # compute usual loss via models
                          _UpperCamelCase       ,		_UpperCamelCase              =						model(**A,       labels=A,       use_cache=A      )[:2]
            else:
                   # compute label smoothed loss
                   _UpperCamelCase              =						model(**A,       use_cache=A      )[0]
                   _UpperCamelCase              =						torch.nn.functional.log_softmax(A,       dim=-1      )
                   _UpperCamelCase       ,		_UpperCamelCase              =						self.loss_fn(A,       A,       self.args.label_smoothing,       ignore_index=self.config.pad_token_id      )
            return loss, logits

     def 	_UpperCamelCase     (      self,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						inputs.pop('''labels'''      )
            _UpperCamelCase       ,		_UpperCamelCase              =						self._compute_loss(A,       A,       A      )
            return loss

     def 	_UpperCamelCase     (      self,       A,       A,       A,       A = None,       ):





            """simple docstring"""

            _UpperCamelCase              =						self._prepare_inputs(A      )

            _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'''],       **A,       )
                   # 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(A,       gen_kwargs['''max_length''']      )

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

            _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(A,       gen_kwargs['''max_length''']      )

            return (loss, logits, labels)



     def 	_UpperCamelCase     (      self,       A,       A      ):





            """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 (
    MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
    TextaTextGenerationPipeline,
    pipeline,
)
from transformers.testing_utils import is_pipeline_test, require_tf, require_torch
from transformers.utils import is_torch_available

from .test_pipelines_common import ANY


if is_torch_available():
       import torch

@is_pipeline_test
class  __A(						unittest.TestCase ):
     __A	  =       MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING
     __A	  =       TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING

     def 	_UpperCamelCase     (      self,       A,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						TextaTextGenerationPipeline(model=A,       tokenizer=A      )
            return generator, ["Something to write", "Something else"]

     def 	_UpperCamelCase     (      self,       A,       A      ):





            """simple docstring"""

            _UpperCamelCase              =						generator('''Something there'''      )
            self.assertEqual(A,       [{'''generated_text''': ANY(A      )}]      )
            # These are encoder decoder, they don't just append to incoming string
            self.assertFalse(outputs[0]['''generated_text'''].startswith('''Something there'''      )      )

            _UpperCamelCase              =						generator(['''This is great !''', '''Something else'''],       num_return_sequences=2,       do_sample=A      )
            self.assertEqual(
                A,       [
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                ],       )

            _UpperCamelCase              =						generator(
                ['''This is great !''', '''Something else'''],       num_return_sequences=2,       batch_size=2,       do_sample=A      )
            self.assertEqual(
                A,       [
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                    [{'''generated_text''': ANY(A      )}, {'''generated_text''': ANY(A      )}],
                ],       )

            with self.assertRaises(A      ):
                   generator(4      )

     @require_torch
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						pipeline('''text2text-generation''',       model='''patrickvonplaten/t5-tiny-random''',       framework='''pt'''      )
            # do_sample=False necessary for reproducibility
            _UpperCamelCase              =						generator('''Something there''',       do_sample=A      )
            self.assertEqual(A,       [{'''generated_text''': ''''''}]      )

            _UpperCamelCase              =						3
            _UpperCamelCase              =						generator(
                '''Something there''',       num_return_sequences=A,       num_beams=A,       )
            _UpperCamelCase              =						[
                {'''generated_text''': '''Beide Beide Beide Beide Beide Beide Beide Beide Beide'''},
                {'''generated_text''': '''Beide Beide Beide Beide Beide Beide Beide Beide'''},
                {'''generated_text''': ''''''},
            ]
            self.assertEqual(A,       A      )

            _UpperCamelCase              =						generator('''This is a test''',       do_sample=A,       num_return_sequences=2,       return_tensors=A      )
            self.assertEqual(
                A,       [
                    {'''generated_token_ids''': ANY(torch.Tensor      )},
                    {'''generated_token_ids''': ANY(torch.Tensor      )},
                ],       )
            _UpperCamelCase              =						generator.model.config.eos_token_id
            _UpperCamelCase              =						'''<pad>'''
            _UpperCamelCase              =						generator(
                ['''This is a test''', '''This is a second test'''],       do_sample=A,       num_return_sequences=2,       batch_size=2,       return_tensors=A,       )
            self.assertEqual(
                A,       [
                    [
                        {'''generated_token_ids''': ANY(torch.Tensor      )},
                        {'''generated_token_ids''': ANY(torch.Tensor      )},
                    ],
                    [
                        {'''generated_token_ids''': ANY(torch.Tensor      )},
                        {'''generated_token_ids''': ANY(torch.Tensor      )},
                    ],
                ],       )



     @require_tf
     def 	_UpperCamelCase     (      self      ):





            """simple docstring"""

            _UpperCamelCase              =						pipeline('''text2text-generation''',       model='''patrickvonplaten/t5-tiny-random''',       framework='''tf'''      )
            # do_sample=False necessary for reproducibility
            _UpperCamelCase              =						generator('''Something there''',       do_sample=A      )
            self.assertEqual(A,       [{'''generated_text''': ''''''}]      )






