





"""simple docstring"""
import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    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.utils.versions import require_version


__UpperCamelCase     :						str						       =  logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

__UpperCamelCase     :						List[str]						       =  {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
__UpperCamelCase     :						int						       =  {
    '''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,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
__UpperCamelCase     :						Any						       =  sorted(arg_to_scheduler.keys())
__UpperCamelCase     :						int						       =  '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''







class       a							(   pl.LightningModule  ):



 def __init__(     self				,       _snake_case				,       _snake_case=None				,       _snake_case="base"				,       _snake_case=None				,       _snake_case=None				,       _snake_case=None				,       **_snake_case				,       ):







  """simple docstring"""






  super().__init__()
  # TODO: move to self.save_hyperparameters()
  # self.save_hyperparameters()
  # can also expand arguments into trainer signature for easier reading

  self.save_hyperparameters(_snake_case     )
  lowerCAmelCase    				= 0
  lowerCAmelCase    				= Path(self.hparams.output_dir     )
  lowerCAmelCase    				= self.hparams.cache_dir if self.hparams.cache_dir else None
  if config is None:
   lowerCAmelCase    				= AutoConfig.from_pretrained(
       self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path				,       **({'num_labels': num_labels} if num_labels is not None else {})				,       cache_dir=_snake_case				,       **_snake_case				,       )
  else:
   lowerCAmelCase    				= config

  lowerCAmelCase    				= ('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
  for p in extra_model_params:
   if getattr(self.hparams				,       _snake_case				,       _snake_case     ):
    assert hasattr(self.config				,       _snake_case     ), F'model config doesn\'t have a `{p}` attribute'
    setattr(self.config				,       _snake_case				,       getattr(self.hparams				,       _snake_case     )     )

  if tokenizer is None:
   lowerCAmelCase    				= AutoTokenizer.from_pretrained(
       self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path				,       cache_dir=_snake_case				,       )
  else:
   lowerCAmelCase    				= tokenizer
  lowerCAmelCase    				= MODEL_MODES[mode]
  if model is None:
   lowerCAmelCase    				= self.model_type.from_pretrained(
       self.hparams.model_name_or_path				,       from_tf=bool('.ckpt' in self.hparams.model_name_or_path     )				,       config=self.config				,       cache_dir=_snake_case				,       )
  else:
   lowerCAmelCase    				= model



 def    UpperCamelCase__  (     self				,       *_snake_case				,       **_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_type.from_pretrained(*_snake_case				,       **_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= arg_to_scheduler[self.hparams.lr_scheduler]
  lowerCAmelCase    				= get_schedule_func(
      self.opt				,       num_warmup_steps=self.hparams.warmup_steps				,       num_training_steps=self.total_steps()     )
  lowerCAmelCase    				= {'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
  return scheduler



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model
  lowerCAmelCase    				= ['bias', 'LayerNorm.weight']
  lowerCAmelCase    				= [
      {
          'params': [
              p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay     )
          ],  # check this named paramters
          'weight_decay': self.hparams.weight_decay,
      },
      {
          'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay     )],
          'weight_decay': 0.0,
      },
  ]
  if self.hparams.adafactor:
   lowerCAmelCase    				= Adafactor(
       _snake_case				,       lr=self.hparams.learning_rate				,       scale_parameter=_snake_case				,       relative_step=_snake_case     )

  else:
   lowerCAmelCase    				= AdamW(
       _snake_case				,       lr=self.hparams.learning_rate				,       eps=self.hparams.adam_epsilon     )
  lowerCAmelCase    				= optimizer

  lowerCAmelCase    				= self.get_lr_scheduler()

  return [optimizer], [scheduler]



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  return self.validation_step(_snake_case				,       _snake_case     )



 def    UpperCamelCase__  (     self				,       _snake_case     ):







  """simple docstring"""






  return self.validation_end(_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= max(1				,       self.hparams.gpus     )  # TODO: consider num_tpu_cores
  lowerCAmelCase    				= self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
  return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs



 def    UpperCamelCase__  (     self				,       _snake_case     ):







  """simple docstring"""






  if stage == "test":
   lowerCAmelCase    				= len(self.test_dataloader().dataset     )
  else:
   lowerCAmelCase    				= self.get_dataloader('train'				,       self.hparams.train_batch_size				,       shuffle=_snake_case     )
   lowerCAmelCase    				= len(self.train_dataloader().dataset     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case = False     ):







  """simple docstring"""






  raise NotImplementedError('You must implement this for your task'     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return self.train_loader



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return self.get_dataloader('dev'				,       self.hparams.eval_batch_size				,       shuffle=_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  return self.get_dataloader('test'				,       self.hparams.eval_batch_size				,       shuffle=_snake_case     )



 def    UpperCamelCase__  (     self				,       _snake_case     ):







  """simple docstring"""






  return os.path.join(
      self.hparams.data_dir				,       'cached_{}_{}_{}'.format(
          _snake_case				,       list(filter(_snake_case				,       self.hparams.model_name_or_path.split('/'     )     )     ).pop()				,       str(self.hparams.max_seq_length     )				,       )				,       )



 @pl.utilities.rank_zero_only
 def    UpperCamelCase__  (     self				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.output_dir.joinpath('best_tfmr'     )
  lowerCAmelCase    				= self.step_count
  self.model.save_pretrained(_snake_case     )
  self.tokenizer.save_pretrained(_snake_case     )



 @staticmethod
 def    UpperCamelCase__  (     _snake_case				,       _snake_case     ):







  """simple docstring"""






  parser.add_argument(
      '--model_name_or_path'				,       default=_snake_case				,       type=_snake_case				,       required=_snake_case				,       help='Path to pretrained model or model identifier from huggingface.co/models'				,       )
  parser.add_argument(
      '--config_name'				,       default=''				,       type=_snake_case				,       help='Pretrained config name or path if not the same as model_name'     )
  parser.add_argument(
      '--tokenizer_name'				,       default=_snake_case				,       type=_snake_case				,       help='Pretrained tokenizer name or path if not the same as model_name'				,       )
  parser.add_argument(
      '--cache_dir'				,       default=str(Path(_snake_case     ).parent / 'test_run' / 'cache'     )				,       type=_snake_case				,       help='Where do you want to store the pre-trained models downloaded from huggingface.co'				,       )
  parser.add_argument(
      '--encoder_layerdrop'				,       type=_snake_case				,       help='Encoder layer dropout probability (Optional). Goes into model.config'				,       )
  parser.add_argument(
      '--decoder_layerdrop'				,       type=_snake_case				,       help='Decoder layer dropout probability (Optional). Goes into model.config'				,       )
  parser.add_argument(
      '--dropout'				,       type=_snake_case				,       help='Dropout probability (Optional). Goes into model.config'				,       )
  parser.add_argument(
      '--attention_dropout'				,       type=_snake_case				,       help='Attention dropout probability (Optional). Goes into model.config'				,       )
  parser.add_argument('--learning_rate'				,       default=5E-5				,       type=_snake_case				,       help='The initial learning rate for Adam.'     )
  parser.add_argument(
      '--lr_scheduler'				,       default='linear'				,       choices=_snake_case				,       metavar=_snake_case				,       type=_snake_case				,       help='Learning rate scheduler'				,       )
  parser.add_argument('--weight_decay'				,       default=0.0				,       type=_snake_case				,       help='Weight decay if we apply some.'     )
  parser.add_argument('--adam_epsilon'				,       default=1E-8				,       type=_snake_case				,       help='Epsilon for Adam optimizer.'     )
  parser.add_argument('--warmup_steps'				,       default=0				,       type=_snake_case				,       help='Linear warmup over warmup_steps.'     )
  parser.add_argument('--num_workers'				,       default=4				,       type=_snake_case				,       help='kwarg passed to DataLoader'     )
  parser.add_argument('--num_train_epochs'				,       dest='max_epochs'				,       default=3				,       type=_snake_case     )
  parser.add_argument('--train_batch_size'				,       default=32				,       type=_snake_case     )
  parser.add_argument('--eval_batch_size'				,       default=32				,       type=_snake_case     )
  parser.add_argument('--adafactor'				,       action='store_true'     )







class       a							(   pl.Callback  ):



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  if (
      trainer.is_global_zero and trainer.global_rank == 0
  ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
   pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.







class       a							(   pl.Callback  ):



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  for name, param in pl_module.model.rag.named_parameters():
   if param.grad is None:
    print(_snake_case     )







class       a							(   pl.Callback  ):



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= trainer.lr_schedulers[0]['scheduler']
  lowerCAmelCase    				= {F'lr_group_{i}': lr for i, lr in enumerate(lr_scheduler.get_lr()     )}
  pl_module.logger.log_metrics(_snake_case     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  rank_zero_info('***** Validation results *****'     )
  lowerCAmelCase    				= trainer.callback_metrics
  # Log results
  for key in sorted(_snake_case     ):
   if key not in ["log", "progress_bar"]:
    rank_zero_info('{} = {}\n'.format(_snake_case				,       str(metrics[key]     )     )     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  rank_zero_info('***** Test results *****'     )
  lowerCAmelCase    				= trainer.callback_metrics
  # Log and save results to file
  lowerCAmelCase    				= os.path.join(pl_module.hparams.output_dir				,       'test_results.txt'     )
  with open(_snake_case				,       'w'     ) as writer:
   for key in sorted(_snake_case     ):
    if key not in ["log", "progress_bar"]:
     rank_zero_info('{} = {}\n'.format(_snake_case				,       str(metrics[key]     )     )     )
     writer.write('{} = {}\n'.format(_snake_case				,       str(metrics[key]     )     )     )


def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    Optional[Any]  ,     _UpperCAmelCase					:    Optional[Any]			):
 #  To allow all pl args uncomment the following line
 #  parser = pl.Trainer.add_argparse_args(parser)
 parser.add_argument(
     '--output_dir'  ,     default=str(Path(_UpperCAmelCase			).parent / 'test_run' / 'model_checkpoints'			)  ,     type=_UpperCAmelCase  ,     help='The output directory where the model predictions and checkpoints will be written.'  ,     )
 parser.add_argument(
     '--fp16'  ,     action='store_true'  ,     help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'  ,     )

 parser.add_argument(
     '--fp16_opt_level'  ,     type=_UpperCAmelCase  ,     default='O2'  ,     help=(
         'For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'
         'See details at https://nvidia.github.io/apex/amp.html'
     )  ,     )
 parser.add_argument('--n_tpu_cores'  ,     dest='tpu_cores'  ,     type=_UpperCAmelCase			)
 parser.add_argument('--max_grad_norm'  ,     dest='gradient_clip_val'  ,     default=1.0  ,     type=_UpperCAmelCase  ,     help='Max gradient norm'			)
 parser.add_argument('--do_train'  ,     action='store_true'  ,     help='Whether to run training.'			)
 parser.add_argument('--do_predict'  ,     action='store_true'  ,     help='Whether to run predictions on the test set.'			)
 parser.add_argument(
     '--gradient_accumulation_steps'  ,     dest='accumulate_grad_batches'  ,     type=_UpperCAmelCase  ,     default=1  ,     help='Number of updates steps to accumulate before performing a backward/update pass.'  ,     )
 parser.add_argument('--seed'  ,     type=_UpperCAmelCase  ,     default=42  ,     help='random seed for initialization'			)
 parser.add_argument(
     '--data_dir'  ,     default=str(Path(_UpperCAmelCase			).parent / 'test_run' / 'dummy-train-data'			)  ,     type=_UpperCAmelCase  ,     help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'  ,     )




def 		_SCREAMING_SNAKE_CASE			(_UpperCAmelCase					:    BaseTransformer  ,     _UpperCAmelCase					:    argparse.Namespace  ,     _UpperCAmelCase					:    List[str]=None  ,     _UpperCAmelCase					:    Union[str, Any]=True  ,     _UpperCAmelCase					:    List[Any]=[]  ,     _UpperCAmelCase					:    List[Any]=None  ,     _UpperCAmelCase					:    List[Any]=None  ,     **_UpperCAmelCase					:    Dict  ,     ):
 pl.seed_everything(args.seed			)

 # init model
 lowerCAmelCase    				= Path(model.hparams.output_dir			)
 odir.mkdir(exist_ok=_UpperCAmelCase			)

 # add custom checkpoints
 if checkpoint_callback is None:
  lowerCAmelCase    				= pl.callbacks.ModelCheckpoint(
      filepath=args.output_dir  ,     prefix='checkpoint'  ,     monitor='val_loss'  ,     mode='min'  ,     save_top_k=1			)
 if early_stopping_callback:
  extra_callbacks.append(_UpperCAmelCase			)
 if logging_callback is None:
  lowerCAmelCase    				= LoggingCallback()

 lowerCAmelCase    				= {}

 if args.fpaa:
  lowerCAmelCase    				= 16

 if args.gpus > 1:
  lowerCAmelCase    				= 'auto'
  lowerCAmelCase    				= 'ddp'

 lowerCAmelCase    				= args.accumulate_grad_batches
 lowerCAmelCase    				= None
 lowerCAmelCase    				= 'auto'

 lowerCAmelCase    				= pl.Trainer.from_argparse_args(
     _UpperCAmelCase  ,     weights_summary=_UpperCAmelCase  ,     callbacks=[logging_callback] + extra_callbacks + [InitCallback()] + [checkpoint_callback]  ,     logger=_UpperCAmelCase  ,     val_check_interval=1  ,     num_sanity_val_steps=2  ,     **_UpperCAmelCase  ,     )

 if args.do_train:
  trainer.fit(_UpperCAmelCase			)

 else:
  print('RAG modeling tests with new set functions successfuly executed!'			)
 return trainer












"""simple docstring"""
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch

       from transformers import (
           OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
           OpenAIGPTConfig,
           OpenAIGPTDoubleHeadsModel,
           OpenAIGPTForSequenceClassification,
           OpenAIGPTLMHeadModel,
           OpenAIGPTModel,
       )







class       a							:



 def __init__(     self				,       _snake_case				,       _snake_case=13				,       _snake_case=7				,       _snake_case=True				,       _snake_case=True				,       _snake_case=True				,       _snake_case=99				,       _snake_case=32				,       _snake_case=5				,       _snake_case=4				,       _snake_case=37				,       _snake_case="gelu"				,       _snake_case=0.1				,       _snake_case=0.1				,       _snake_case=5_12				,       _snake_case=16				,       _snake_case=2				,       _snake_case=0.02				,       _snake_case=3				,       _snake_case=4				,       _snake_case=None				,       ):







  """simple docstring"""






  lowerCAmelCase    				= parent
  lowerCAmelCase    				= batch_size
  lowerCAmelCase    				= seq_length
  lowerCAmelCase    				= is_training
  lowerCAmelCase    				= use_token_type_ids
  lowerCAmelCase    				= use_labels
  lowerCAmelCase    				= vocab_size
  lowerCAmelCase    				= hidden_size
  lowerCAmelCase    				= num_hidden_layers
  lowerCAmelCase    				= num_attention_heads
  lowerCAmelCase    				= intermediate_size
  lowerCAmelCase    				= hidden_act
  lowerCAmelCase    				= hidden_dropout_prob
  lowerCAmelCase    				= attention_probs_dropout_prob
  lowerCAmelCase    				= max_position_embeddings
  lowerCAmelCase    				= type_vocab_size
  lowerCAmelCase    				= type_sequence_label_size
  lowerCAmelCase    				= initializer_range
  lowerCAmelCase    				= num_labels
  lowerCAmelCase    				= num_choices
  lowerCAmelCase    				= scope
  lowerCAmelCase    				= self.vocab_size - 1



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size     )

  lowerCAmelCase    				= None
  if self.use_token_type_ids:
   lowerCAmelCase    				= ids_tensor([self.batch_size, self.seq_length]				,       self.type_vocab_size     )

  lowerCAmelCase    				= None
  lowerCAmelCase    				= None
  lowerCAmelCase    				= None
  if self.use_labels:
   lowerCAmelCase    				= ids_tensor([self.batch_size]				,       self.type_sequence_label_size     )
   lowerCAmelCase    				= ids_tensor([self.batch_size, self.seq_length]				,       self.num_labels     )
   lowerCAmelCase    				= ids_tensor([self.batch_size]				,       self.num_choices     )

  lowerCAmelCase    				= OpenAIGPTConfig(
      vocab_size=self.vocab_size				,       n_embd=self.hidden_size				,       n_layer=self.num_hidden_layers				,       n_head=self.num_attention_heads				,       n_positions=self.max_position_embeddings				,       pad_token_id=self.pad_token_id				,       )

  lowerCAmelCase    				= ids_tensor([self.num_hidden_layers, self.num_attention_heads]				,       2     )

  return (
      config,
      input_ids,
      head_mask,
      token_type_ids,
      sequence_labels,
      token_labels,
      choice_labels,
  )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModel(config=_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       head_mask=_snake_case     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case     )
  lowerCAmelCase    				= model(_snake_case     )

  self.parent.assertEqual(result.last_hidden_state.shape				,       (self.batch_size, self.seq_length, self.hidden_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTDoubleHeadsModel(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.loss.shape				,       ()     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.seq_length, self.vocab_size)     )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       *_snake_case     ):







  """simple docstring"""






  lowerCAmelCase    				= self.num_labels
  lowerCAmelCase    				= OpenAIGPTForSequenceClassification(_snake_case     )
  model.to(_snake_case     )
  model.eval()

  lowerCAmelCase    				= ids_tensor([self.batch_size]				,       self.type_sequence_label_size     )
  lowerCAmelCase    				= model(_snake_case				,       token_type_ids=_snake_case				,       labels=_snake_case     )
  self.parent.assertEqual(result.logits.shape				,       (self.batch_size, self.num_labels)     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.prepare_config_and_inputs()
  (
      (
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,(
      lowerCAmelCase
  )				,
  )    				= config_and_inputs
  lowerCAmelCase    				= {
      'input_ids': input_ids,
      'token_type_ids': token_type_ids,
      'head_mask': head_mask,
  }

  return config, inputs_dict







@require_torch
class       a							(   a__   , a__   , a__   , unittest.TestCase  ):
 snake_case__             =       (
     (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
     if is_torch_available()
     else ()
 )
 snake_case__             =       (
     (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
 )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
 snake_case__             =       (
     {
         '''feature-extraction''': OpenAIGPTModel,
         '''text-classification''': OpenAIGPTForSequenceClassification,
         '''text-generation''': OpenAIGPTLMHeadModel,
         '''zero-shot''': OpenAIGPTForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case				,       _snake_case     ):







  """simple docstring"""






  if pipeline_test_casse_name == "ZeroShotClassificationPipelineTests":
   # Get `tokenizer does not have a padding token` error for both fast/slow tokenizers.
   # `OpenAIGPTConfig` was never used in pipeline tests, either because of a missing checkpoint or because a
   # tiny config could not be created.
   return True

  return False



 def    UpperCamelCase__  (     self				,       _snake_case				,       _snake_case				,       _snake_case=False     ):







  """simple docstring"""






  lowerCAmelCase    				= super()._prepare_for_class(_snake_case				,       _snake_case				,       return_labels=_snake_case     )

  if return_labels:
   if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= inputs_dict['labels']
    lowerCAmelCase    				= torch.zeros(
        (self.model_tester.batch_size, self.model_tester.num_choices)				,       dtype=torch.long				,       device=_snake_case				,       )
    lowerCAmelCase    				= torch.zeros(
        self.model_tester.batch_size				,       dtype=torch.long				,       device=_snake_case     )
  return inputs_dict



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTModelTester(self     )
  lowerCAmelCase    				= ConfigTester(self				,       config_class=_snake_case				,       n_embd=37     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  self.config_tester.run_common_tests()



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_double_lm_head_model(*_snake_case     )



 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*_snake_case     )



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   lowerCAmelCase    				= OpenAIGPTModel.from_pretrained(_snake_case     )
   self.assertIsNotNone(_snake_case     )







@require_torch
class       a							(   unittest.TestCase  ):



 @slow
 def    UpperCamelCase__  (     self     ):







  """simple docstring"""






  lowerCAmelCase    				= OpenAIGPTLMHeadModel.from_pretrained('openai-gpt'     )
  model.to(_snake_case     )
  lowerCAmelCase    				= torch.tensor([[4_81, 47_35, 5_44]]				,       dtype=torch.long				,       device=_snake_case     )  # the president is
  lowerCAmelCase    				= [
      4_81,
      47_35,
      5_44,
      2_46,
      9_63,
      8_70,
      7_62,
      2_39,
      2_44,
      4_04_77,
      2_44,
      2_49,
      7_19,
      8_81,
      4_87,
      5_44,
      2_40,
      2_44,
      6_03,
      4_81,
  ]  # the president is a very good man. " \n " i\'m sure he is, " said the

  lowerCAmelCase    				= model.generate(_snake_case				,       do_sample=_snake_case     )
  self.assertListEqual(output_ids[0].tolist()				,       _snake_case     )





