




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

lowerCamelCase	:	Tuple				=   {
    '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 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





      def __init__(				self    ,				A=None    ,				A=None    ,				*A    ,				**A     )		->    Union[str, Any]:
         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__}"""
            )
            snake_case				:    Optional[int]           =	self.model.config
         else:
            snake_case				:    Dict           =	config

         snake_case				:    Tuple           =	data_args
         snake_case				:    str           =	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:
            snake_case				:    Optional[Any]           =	torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id     )
         else:
            # dynamically import label_smoothed_nll_loss
            from utils import label_smoothed_nll_loss

            snake_case				:    Dict           =	label_smoothed_nll_loss

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

         if self.lr_scheduler is None:
            snake_case				:    Optional[Any]           =	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     )		->    Optional[int]:
         snake_case				:    str           =	arg_to_scheduler[self.args.lr_scheduler]
         if self.args.lr_scheduler == "constant":
            snake_case				:    Optional[Any]           =	schedule_func(self.optimizer     )
         elif self.args.lr_scheduler == "constant_w_warmup":
            snake_case				:    int           =	schedule_func(self.optimizer    ,				num_warmup_steps=self.args.warmup_steps     )
         else:
            snake_case				:    Tuple           =	schedule_func(
                self.optimizer    ,				num_warmup_steps=self.args.warmup_steps    ,				num_training_steps=A     )
         return scheduler

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

      def      UpperCAmelCase				(				self    ,				A    ,				A     )		->    Union[str, Any]:
         snake_case				:    List[str]           =	inputs.pop("""labels"""     )
         snake_case     , snake_case				:    Any           =	self._compute_loss(A    ,				A    ,				A     )
         return loss

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A = None    ,				)		->    Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
         snake_case				:    Optional[int]           =	self._prepare_inputs(A     )

         snake_case				:    Dict           =	{
             """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:
            snake_case				:    List[str]           =	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"]:
               snake_case				:    List[str]           =	self._pad_tensors_to_max_len(A    ,				gen_kwargs["""max_length"""]     )

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

         snake_case				:    Union[str, Any]           =	loss.mean().detach()
         if self.args.prediction_loss_only:
            return (loss, None, None)

         snake_case				:    Dict           =	generated_tokens if self.args.predict_with_generate else logits

         if labels.shape[-1] < gen_kwargs["max_length"]:
            snake_case				:    Optional[int]           =	self._pad_tensors_to_max_len(A    ,				gen_kwargs["""max_length"""]     )

         return (loss, logits, labels)

      def      UpperCAmelCase				(				self    ,				A    ,				A     )		->    List[str]:
         # If PAD token is not defined at least EOS token has to be defined
         snake_case				:    Tuple           =	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}"""     )

         snake_case				:    Dict           =	pad_token_id * torch.ones(
             (tensor.shape[0], max_length)    ,				dtype=tensor.dtype    ,				device=tensor.device     )
         snake_case				:    Tuple           =	tensor
         return padded_tensor











from argparse import ArgumentParser, Namespace

from ..utils import logging
from . import BaseTransformersCLICommand




def 						SCREAMING_SNAKE_CASE__ (   lowercase					)  ->       List[str]:
   return ConvertCommand(
       args.model_type						,args.tf_checkpoint						,args.pytorch_dump_output						,args.config						,args.finetuning_task_name					)


lowerCamelCase	:	Tuple				=   '\ntransformers can only be used from the commandline to convert TensorFlow models in PyTorch, In that case, it requires\nTensorFlow to be installed. Please see https://www.tensorflow.org/install/ for installation instructions.\n'




class 						__lowercase		(UpperCamelCase__							):




      """simple docstring"""





      @staticmethod
      def      UpperCAmelCase				(				A     )		->    List[str]:
         snake_case				:    Union[str, Any]           =	parser.add_parser(
             """convert"""    ,				help="""CLI tool to run convert model from original author checkpoints to Transformers PyTorch checkpoints."""    ,				)
         train_parser.add_argument("""--model_type"""    ,				type=A    ,				required=A    ,				help="""Model's type."""     )
         train_parser.add_argument(
             """--tf_checkpoint"""    ,				type=A    ,				required=A    ,				help="""TensorFlow checkpoint path or folder."""     )
         train_parser.add_argument(
             """--pytorch_dump_output"""    ,				type=A    ,				required=A    ,				help="""Path to the PyTorch saved model output."""     )
         train_parser.add_argument("""--config"""    ,				type=A    ,				default=""""""    ,				help="""Configuration file path or folder."""     )
         train_parser.add_argument(
             """--finetuning_task_name"""    ,				type=A    ,				default=A    ,				help="""Optional fine-tuning task name if the TF model was a finetuned model."""    ,				)
         train_parser.set_defaults(func=A     )

      def __init__(				self    ,				A    ,				A    ,				A    ,				A    ,				A    ,				*A    ,				)		->    List[Any]:
         snake_case				:    Any           =	logging.get_logger("""transformers-cli/converting"""     )

         self._logger.info(f"""Loading model {model_type}"""     )
         snake_case				:    int           =	model_type
         snake_case				:    Any           =	tf_checkpoint
         snake_case				:    int           =	pytorch_dump_output
         snake_case				:    List[str]           =	config
         snake_case				:    Tuple           =	finetuning_task_name

      def      UpperCAmelCase				(				self     )		->    int:
         if self._model_type == "albert":
            try:
               from ..models.albert.convert_albert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "bert":
            try:
               from ..models.bert.convert_bert_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "funnel":
            try:
               from ..models.funnel.convert_funnel_original_tf_checkpoint_to_pytorch import (
                   convert_tf_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "t5":
            try:
               from ..models.ta.convert_ta_original_tf_checkpoint_to_pytorch import convert_tf_checkpoint_to_pytorch
            except ImportError:
               raise ImportError(A     )

            convert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "gpt":
            from ..models.openai.convert_openai_original_tf_checkpoint_to_pytorch import (
                convert_openai_checkpoint_to_pytorch,
            )

            convert_openai_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "transfo_xl":
            try:
               from ..models.transfo_xl.convert_transfo_xl_original_tf_checkpoint_to_pytorch import (
                   convert_transfo_xl_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            if "ckpt" in self._tf_checkpoint.lower():
               snake_case				:    List[Any]           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            else:
               snake_case				:    Tuple           =	self._tf_checkpoint
               snake_case				:    Tuple           =	""""""
            convert_transfo_xl_checkpoint_to_pytorch(
                A    ,				self._config    ,				self._pytorch_dump_output    ,				A     )
         elif self._model_type == "gpt2":
            try:
               from ..models.gpta.convert_gpta_original_tf_checkpoint_to_pytorch import (
                   convert_gpta_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_gpta_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         elif self._model_type == "xlnet":
            try:
               from ..models.xlnet.convert_xlnet_original_tf_checkpoint_to_pytorch import (
                   convert_xlnet_checkpoint_to_pytorch,
               )
            except ImportError:
               raise ImportError(A     )

            convert_xlnet_checkpoint_to_pytorch(
                self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output    ,				self._finetuning_task_name     )
         elif self._model_type == "xlm":
            from ..models.xlm.convert_xlm_original_pytorch_checkpoint_to_pytorch import (
                convert_xlm_checkpoint_to_pytorch,
            )

            convert_xlm_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "lxmert":
            from ..models.lxmert.convert_lxmert_original_tf_checkpoint_to_pytorch import (
                convert_lxmert_checkpoint_to_pytorch,
            )

            convert_lxmert_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._pytorch_dump_output     )
         elif self._model_type == "rembert":
            from ..models.rembert.convert_rembert_tf_checkpoint_to_pytorch import (
                convert_rembert_tf_checkpoint_to_pytorch,
            )

            convert_rembert_tf_checkpoint_to_pytorch(self._tf_checkpoint    ,				self._config    ,				self._pytorch_dump_output     )
         else:
            raise ValueError(
                """--model_type should be selected in the list [bert, gpt, gpt2, t5, transfo_xl, xlnet, xlm, lxmert]"""     )





