


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     _lowerCamelCase					(   lowerCamelCase_     ):




       """simple docstring"""



       def __init__( self							,			_SCREAMING_SNAKE_CASE=None							,			_SCREAMING_SNAKE_CASE=None							,			*_SCREAMING_SNAKE_CASE							,			**_SCREAMING_SNAKE_CASE						)->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__}'''
           )
           A_ :							Any							=   self.model.config
         else:
           A_ :							Dict							=   config

         A_ :							Optional[Any]							=   data_args
         A_ :							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:
           A_ :							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

           A_ :							Optional[Any]							=   label_smoothed_nll_loss



       def 			_snake_case			( self							,			_SCREAMING_SNAKE_CASE						)->Dict:




         '''simple docstring'''





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

         if self.lr_scheduler is None:
           A_ :							int							=   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							,			_SCREAMING_SNAKE_CASE						)->str:




         '''simple docstring'''





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



       def 			_snake_case			( self						)->int:




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




         '''simple docstring'''





         A_ :							List[Any]							=   inputs.pop('''labels'''						)
         A_ :							Optional[Any]							=   self._compute_loss(snake_case__							,			snake_case__							,			snake_case__						)
         return loss



       def 			_snake_case			( self							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE = None							,			)->str:




         '''simple docstring'''





         A_ :							Tuple							=   self._prepare_inputs(snake_case__						)

         A_ :							List[str]							=   {
             '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:
           A_ :							str							=   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"]:
             A_ :							Tuple							=   self._pad_tensors_to_max_len(snake_case__							,			gen_kwargs['''max_length''']						)

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

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

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

         if labels.shape[-1] < gen_kwargs["max_length"]:
           A_ :							Tuple							=   self._pad_tensors_to_max_len(snake_case__							,			gen_kwargs['''max_length''']						)

         return (loss, logits, labels)







       def 			_snake_case			( self							,			_SCREAMING_SNAKE_CASE							,			_SCREAMING_SNAKE_CASE						)->List[Any]:




         '''simple docstring'''





         A_ :							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}'''						)

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










from __future__ import annotations

import unittest

from transformers import XGLMConfig, XGLMTokenizer, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers.models.xglm.modeling_tf_xglm import (
						    TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST,
						    TFXGLMForCausalLM,
						    TFXGLMModel,
						)







@require_tf
class 				lowerCAmelCase_       :
				__a     :       List[Any]						=     XGLMConfig
				__a     :       List[Any]						=     {}
				__a     :       List[str]						=     "gelu"






				def __init__(			self      ,snake_case__      ,snake_case__=14      ,snake_case__=7      ,snake_case__=True      ,snake_case__=True      ,snake_case__=True      ,snake_case__=99      ,snake_case__=32      ,snake_case__=2      ,snake_case__=4      ,snake_case__=37      ,snake_case__="gelu"      ,snake_case__=0.1      ,snake_case__=0.1      ,snake_case__=512      ,snake_case__=0.02      ,):
							SCREAMING_SNAKE_CASE_   : int					  =      parent
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      batch_size
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      seq_length
							SCREAMING_SNAKE_CASE_   : List[str]					  =      is_training
							SCREAMING_SNAKE_CASE_   : int					  =      use_input_mask
							SCREAMING_SNAKE_CASE_   : Tuple					  =      use_labels
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      vocab_size
							SCREAMING_SNAKE_CASE_   : int					  =      d_model
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      num_hidden_layers
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      num_attention_heads
							SCREAMING_SNAKE_CASE_   : List[str]					  =      ffn_dim
							SCREAMING_SNAKE_CASE_   : Any					  =      activation_function
							SCREAMING_SNAKE_CASE_   : Dict					  =      activation_dropout
							SCREAMING_SNAKE_CASE_   : int					  =      attention_dropout
							SCREAMING_SNAKE_CASE_   : str					  =      max_position_embeddings
							SCREAMING_SNAKE_CASE_   : Any					  =      initializer_range
							SCREAMING_SNAKE_CASE_   : int					  =      None
							SCREAMING_SNAKE_CASE_   : str					  =      0
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      2
							SCREAMING_SNAKE_CASE_   : Any					  =      1






				def 		snake_case	(			self ):
							return XGLMConfig.from_pretrained('facebook/xglm-564M' )






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      tf.clip_by_value(
							    ids_tensor([self.batch_size, self.seq_length]      ,self.vocab_size )      ,clip_value_min=0      ,clip_value_max=3 )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      None
							if self.use_input_mask:
										SCREAMING_SNAKE_CASE_   : List[str]					  =      random_attention_mask([self.batch_size, self.seq_length] )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      self.get_config()

							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      floats_tensor([self.num_hidden_layers, self.num_attention_heads]      ,2 )

							return (
							    config,
							    input_ids,
							    input_mask,
							    head_mask,
							)






				def 		snake_case	(			self ):
							return XGLMConfig(
							    vocab_size=self.vocab_size      ,d_model=self.hidden_size      ,num_layers=self.num_hidden_layers      ,attention_heads=self.num_attention_heads      ,ffn_dim=self.ffn_dim      ,activation_function=self.activation_function      ,activation_dropout=self.activation_dropout      ,attention_dropout=self.attention_dropout      ,max_position_embeddings=self.max_position_embeddings      ,initializer_range=self.initializer_range      ,use_cache=snake_case__      ,bos_token_id=self.bos_token_id      ,eos_token_id=self.eos_token_id      ,pad_token_id=self.pad_token_id      ,return_dict=snake_case__      ,)






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : str					  =      self.prepare_config_and_inputs()

							(
							    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    (
							    SCREAMING_SNAKE_CASE_
							)  ,    
							)   : List[Any]					  =      config_and_inputs

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      {
							    'input_ids': input_ids,
							    'head_mask': head_mask,
							}

							return config, inputs_dict







@require_tf
class 				lowerCAmelCase_       (					lowerCamelCase_						,     lowerCamelCase_						,     unittest.TestCase							):
				__a     :       Dict						=     (TFXGLMModel, TFXGLMForCausalLM) if is_tf_available() else ()
				__a     :       Dict						=     (TFXGLMForCausalLM,) if is_tf_available() else ()
				__a     :       Union[str, Any]						=     (
				    {"feature-extraction": TFXGLMModel, "text-generation": TFXGLMForCausalLM} if is_tf_available() else {}
				)
				__a     :       Dict						=     False
				__a     :       Union[str, Any]						=     False
				__a     :       List[str]						=     False






				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : int					  =      TFXGLMModelTester(self )
							SCREAMING_SNAKE_CASE_   : List[Any]					  =      ConfigTester(self      ,config_class=snake_case__      ,n_embd=37 )






				def 		snake_case	(			self ):
							self.config_tester.run_common_tests()






				@slow
				def 		snake_case	(			self ):
							for model_name in TF_XGLM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
										SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      TFXGLMModel.from_pretrained(snake_case__ )
										self.assertIsNotNone(snake_case__ )






				@unittest.skip(reason='Currently, model embeddings are going to undergo a major refactor.' )
				def 		snake_case	(			self ):
							super().test_resize_token_embeddings()







@require_tf
class 				lowerCAmelCase_       (					unittest.TestCase							):






				@slow
				def 		snake_case	(			self      ,snake_case__=True ):
							SCREAMING_SNAKE_CASE_   : Dict					  =      TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M' )
							SCREAMING_SNAKE_CASE_   : Any					  =      tf.convert_to_tensor([[2, 268, 9865]]      ,dtype=tf.intaa )  # The dog
							# </s> The dog is a very friendly dog. He is very affectionate and loves to play with other
							# fmt: off
							SCREAMING_SNAKE_CASE_   : Tuple					  =      [2, 268, 9865, 67, 11, 1988, 57252, 9865, 5, 984, 67, 1988, 213838, 1658, 53, 70446, 33, 6657, 278, 1581]
							# fmt: on
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      model.generate(snake_case__      ,do_sample=snake_case__      ,num_beams=1 )
							if verify_outputs:
										self.assertListEqual(output_ids[0].numpy().tolist()      ,snake_case__ )






				@slow
				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Tuple					  =      XGLMTokenizer.from_pretrained('facebook/xglm-564M' )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M' )

							tf.random.set_seed(0 )
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      tokenizer('Today is a nice day and'      ,return_tensors='tf' )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      tokenized.input_ids
							# forces the generation to happen on CPU, to avoid GPU-related quirks (and assure same output regardless of the available devices)
							with tf.device(':/CPU:0' ):
										SCREAMING_SNAKE_CASE_   : List[Any]					  =      model.generate(snake_case__      ,do_sample=snake_case__      ,seed=[7, 0] )
							SCREAMING_SNAKE_CASE_   : Optional[int]					  =      tokenizer.decode(output_ids[0]      ,skip_special_tokens=snake_case__ )

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      (
							    'Today is a nice day and warm evening here over Southern Alberta!! Today when they closed schools due'
							)
							self.assertEqual(snake_case__      ,snake_case__ )






				@slow
				def 		snake_case	(			self ):
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      TFXGLMForCausalLM.from_pretrained('facebook/xglm-564M' )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      XGLMTokenizer.from_pretrained('facebook/xglm-564M' )

							SCREAMING_SNAKE_CASE_   : List[Any]					  =      'left'

							# use different length sentences to test batching
							SCREAMING_SNAKE_CASE_   : Union[str, Any]					  =      [
							    'This is an extremelly long sentence that only exists to test the ability of the model to cope with '
							    'left-padding, such as in batched generation. The output for the sequence below should be the same '
							    'regardless of whether left padding is applied or not. When',
							    'Hello, my dog is a little',
							]

							SCREAMING_SNAKE_CASE_   : Any					  =      tokenizer(snake_case__      ,return_tensors='tf'      ,padding=snake_case__ )
							SCREAMING_SNAKE_CASE_   : int					  =      inputs['input_ids']

							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      model.generate(input_ids=snake_case__      ,attention_mask=inputs['attention_mask']      ,max_new_tokens=12 )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      tokenizer(sentences[0]      ,return_tensors='tf' ).input_ids
							SCREAMING_SNAKE_CASE_   : Any					  =      model.generate(input_ids=snake_case__      ,max_new_tokens=12 )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      tokenizer(sentences[1]      ,return_tensors='tf' ).input_ids
							SCREAMING_SNAKE_CASE_   : Any					  =      model.generate(input_ids=snake_case__      ,max_new_tokens=12 )

							SCREAMING_SNAKE_CASE_   : List[str]					  =      tokenizer.batch_decode(snake_case__      ,skip_special_tokens=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Optional[Any]					  =      tokenizer.decode(output_non_padded[0]      ,skip_special_tokens=snake_case__ )
							SCREAMING_SNAKE_CASE_   : Tuple					  =      tokenizer.decode(output_padded[0]      ,skip_special_tokens=snake_case__ )

							SCREAMING_SNAKE_CASE_   : str					  =      [
							    'This is an extremelly long sentence that only exists to test the ability of the model to cope with '
							    'left-padding, such as in batched generation. The output for the sequence below should be the same '
							    'regardless of whether left padding is applied or not. When left padding is applied, the sequence will be '
							    'a single',
							    'Hello, my dog is a little bit of a shy one, but he is very friendly',
							]
							self.assertListEqual(snake_case__      ,snake_case__ )
							self.assertListEqual(snake_case__      ,[non_padded_sentence, padded_sentence] )



