


"""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


UpperCAmelCase      		=			logging.get_logger(__name__)

UpperCAmelCase      		=			{
    '''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 					__magic_name__		(			__UpperCAmelCase     ):






       def __init__(   self  :    Union[str, Any]  ,     snake_case__  :    Tuple=None  ,     snake_case__  :    List[str]=None  ,     *snake_case__  :    Tuple  ,     **snake_case__  :    Union[str, Any]				):

              '''simple docstring'''






              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__}"""
                     )
                     lowercase		:Optional[Any]       							=				self.model.config
              else:
                     lowercase		:Optional[int]       							=				config

              lowercase		:Optional[Any]       							=				data_args
              lowercase		:Union[str, Any]       							=				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:
                     lowercase		:List[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

                     lowercase		:Union[str, Any]       							=				label_smoothed_nll_loss






       def  __snake_case						(   self  :    Union[str, Any]  ,     snake_case__  :    int				):

              '''simple docstring'''






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

              if self.lr_scheduler is None:
                     lowercase		:Union[str, Any]       							=				self._get_lr_scheduler(snake_case__				)
              else:  # ignoring --lr_scheduler
                     logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''				)






       def  __snake_case						(   self  :    str  ,     snake_case__  :    str				):

              '''simple docstring'''






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






       def  __snake_case						(   self  :    Tuple				):

              '''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  __snake_case						(   self  :    List[Any]  ,     snake_case__  :    Tuple  ,     snake_case__  :    str  ,     snake_case__  :    Any				):

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






       def  __snake_case						(   self  :    List[Any]  ,     snake_case__  :    Any  ,     snake_case__  :    int				):

              '''simple docstring'''






              lowercase		:Optional[int]       							=				inputs.pop('''labels'''				)
              lowercase							,  lowercase		:int       							=				self._compute_loss(snake_case__  ,     snake_case__  ,     snake_case__				)
              return loss






       def  __snake_case						(   self  :    Optional[int]  ,     snake_case__  :    nn.Module  ,     snake_case__  :    Dict[str, Union[torch.Tensor, Any]]  ,     snake_case__  :    bool  ,     snake_case__  :    Optional[List[str]] = None  ,     ):

              '''simple docstring'''






              lowercase		:str       							=				self._prepare_inputs(snake_case__				)

              lowercase		: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:
                     lowercase		:Any       							=				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"]:
                            lowercase		:Optional[Any]       							=				self._pad_tensors_to_max_len(snake_case__  ,     gen_kwargs['''max_length''']				)

              lowercase		:List[Any]       							=				inputs.pop('''labels'''				)
              with torch.no_grad():
                     # compute loss on predict data
                     lowercase							,  lowercase		:List[str]       							=				self._compute_loss(snake_case__  ,     snake_case__  ,     snake_case__				)

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

              lowercase		:Any       							=				generated_tokens if self.args.predict_with_generate else logits

              if labels.shape[-1] < gen_kwargs["max_length"]:
                     lowercase		:Union[str, Any]       							=				self._pad_tensors_to_max_len(snake_case__  ,     gen_kwargs['''max_length''']				)

              return (loss, logits, labels)







       def  __snake_case						(   self  :    Union[str, Any]  ,     snake_case__  :    Union[str, Any]  ,     snake_case__  :    List[str]				):

              '''simple docstring'''






              lowercase		:Union[str, Any]       							=				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}"""				)

              lowercase		:int       							=				pad_token_id * torch.ones(
                  (tensor.shape[0], max_length)  ,     dtype=tensor.dtype  ,     device=tensor.device				)
              lowercase		:Optional[Any]       							=				tensor
              return padded_tensor





"""simple docstring"""



# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import re

from ..utils import cached_file


# docstyle-ignore
UpperCAmelCase      		=			'''
Human: <<task>>

Assistant: '''


UpperCAmelCase      		=			'''huggingface-tools/default-prompts'''
UpperCAmelCase      		=			{'''chat''': '''chat_prompt_template.txt''', '''run''': '''run_prompt_template.txt'''}






def       lowerCamelCase					(a_						:int				,				a_						:str				,				a_						:Dict="run")     ->							Optional[Any]:
       if prompt_or_repo_id is None:
              lowercase		:Tuple       							=				DEFAULT_PROMPTS_REPO

       # prompt is considered a repo ID when it does not contain any kind of space
       if re.search('''\\s'''				,				a_) is not None:
              return prompt_or_repo_id

       lowercase		:List[str]       							=				cached_file(
           a_				,				PROMPT_FILES[mode]				,				repo_type='''dataset'''				,				user_agent={'''agent''': agent_name})
       with open(a_				,				'''r'''				,				encoding='''utf-8''') as f:
              return f.read()

