



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


a_      :      Union[str, Any]    =     logging.getLogger(__name__)

require_version('pytorch_lightning>=1.0.4')

a_      :      List[Any]    =     {
    '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
a_      :      Optional[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
}
a_      :      Tuple    =     sorted(arg_to_scheduler.keys())
a_      :      int    =     '{' + ', '.join(arg_to_scheduler_choices) + '}'
class    __UpperCamelCase			(       pl.LightningModule				):

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE="base"		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None		,				SCREAMING_SNAKE_CASE=None		,				**SCREAMING_SNAKE_CASE		,				) ->			List[str]:
  super().__init__()
  # TODO: move to self.save_hyperparameters()
  # self.save_hyperparameters()
  # can also expand arguments into trainer signature for easier reading

  self.save_hyperparameters(SCREAMING_SNAKE_CASE      )
  a__												= 0
  a__												= Path(self.hparams.output_dir      )
  a__												= self.hparams.cache_dir if self.hparams.cache_dir else None
  if config is None:
   a__												= 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=SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE		,				)
  else:
   a__												= config

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

  if tokenizer is None:
   a__												= AutoTokenizer.from_pretrained(
       self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path		,				cache_dir=SCREAMING_SNAKE_CASE		,				)
  else:
   a__												= tokenizer
  a__												= MODEL_MODES[mode]
  if model is None:
   a__												= 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=SCREAMING_SNAKE_CASE		,				)
  else:
   a__												= model

 def 						_UpperCAmelCase       (	self		,				*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      ) ->			Optional[int]:
  a__												= self.model_type.from_pretrained(*SCREAMING_SNAKE_CASE		,				**SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= arg_to_scheduler[self.hparams.lr_scheduler]
  a__												= get_schedule_func(
      self.opt		,				num_warmup_steps=self.hparams.warmup_steps		,				num_training_steps=self.total_steps()      )
  a__												= {'''scheduler''': scheduler, '''interval''': '''step''', '''frequency''': 1}
  return scheduler

 def 						_UpperCAmelCase       (	self      ) ->			Tuple:
  a__												= self.model
  a__												= ['''bias''', '''LayerNorm.weight''']
  a__												= [
      {
          '''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:
   a__												= Adafactor(
       SCREAMING_SNAKE_CASE		,				lr=self.hparams.learning_rate		,				scale_parameter=SCREAMING_SNAKE_CASE		,				relative_step=SCREAMING_SNAKE_CASE      )

  else:
   a__												= AdamW(
       SCREAMING_SNAKE_CASE		,				lr=self.hparams.learning_rate		,				eps=self.hparams.adam_epsilon      )
  a__												= optimizer

  a__												= self.get_lr_scheduler()

  return [optimizer], [scheduler]

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Union[str, Any]:
  return self.validation_step(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  return self.validation_end(SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= max(1		,				self.hparams.gpus      )  # TODO: consider num_tpu_cores
  a__												= 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		,				SCREAMING_SNAKE_CASE      ) ->			Any:
  if stage == "test":
   a__												= len(self.test_dataloader().dataset      )
  else:
   a__												= self.get_dataloader('''train'''		,				self.hparams.train_batch_size		,				shuffle=SCREAMING_SNAKE_CASE      )
   a__												= len(self.train_dataloader().dataset      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE = False      ) ->			str:
  raise NotImplementedError('''You must implement this for your task'''      )

 def 						_UpperCAmelCase       (	self      ) ->			Dict:
  return self.train_loader

 def 						_UpperCAmelCase       (	self      ) ->			List[Any]:
  return self.get_dataloader('''dev'''		,				self.hparams.eval_batch_size		,				shuffle=SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  return self.get_dataloader('''test'''		,				self.hparams.eval_batch_size		,				shuffle=SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			str:
  return os.path.join(
      self.hparams.data_dir		,				'''cached_{}_{}_{}'''.format(
          SCREAMING_SNAKE_CASE		,				list(filter(SCREAMING_SNAKE_CASE		,				self.hparams.model_name_or_path.split('''/'''      )      )      ).pop()		,				str(self.hparams.max_seq_length      )		,				)		,				)

 @pl.utilities.rank_zero_only
 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE      ) ->			None:
  a__												= self.output_dir.joinpath('''best_tfmr'''      )
  a__												= self.step_count
  self.model.save_pretrained(SCREAMING_SNAKE_CASE      )
  self.tokenizer.save_pretrained(SCREAMING_SNAKE_CASE      )



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





class    __UpperCamelCase			(       pl.Callback				):

 """simple docstring"""



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			List[Any]:
  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    __UpperCamelCase			(       pl.Callback				):

 """simple docstring"""



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			List[Any]:
  # print(pl_module.model.rag)
  for name, param in pl_module.model.rag.named_parameters():
   if param.grad is None:
    print(SCREAMING_SNAKE_CASE      )





class    __UpperCamelCase			(       pl.Callback				):

 """simple docstring"""



 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Tuple:
  a__												= trainer.lr_schedulers[0]['''scheduler''']
  a__												= {f"lr_group_{i}": lr for i, lr in enumerate(lr_scheduler.get_lr()      )}
  pl_module.logger.log_metrics(SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			List[Any]:
  rank_zero_info('''***** Validation results *****'''      )
  a__												= trainer.callback_metrics
  # Log results
  for key in sorted(SCREAMING_SNAKE_CASE      ):
   if key not in ["log", "progress_bar"]:
    rank_zero_info('''{} = {}\n'''.format(SCREAMING_SNAKE_CASE		,				str(metrics[key]      )      )      )



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






def     __a					(		__UpperCAmelCase							,  __UpperCAmelCase       ):
 #  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     __a					(		__UpperCAmelCase							,  __UpperCAmelCase							,  __UpperCAmelCase=None							,  __UpperCAmelCase=True							,  __UpperCAmelCase=[]							,  __UpperCAmelCase=None							,  __UpperCAmelCase=None							,  **__UpperCAmelCase							,  ):
 pl.seed_everything(args.seed       )

 # init model
 a__												= Path(model.hparams.output_dir       )
 odir.mkdir(exist_ok=__UpperCAmelCase       )

 # add custom checkpoints
 if checkpoint_callback is None:
  a__												= 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:
  a__												= LoggingCallback()

 a__												= {}

 if args.fpaa:
  a__												= 16

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

 a__												= args.accumulate_grad_batches
 a__												= None
 a__												= '''auto'''

 a__												= 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










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    __UpperCamelCase			:

 """simple docstring"""



 def __init__(	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=1_3		,				SCREAMING_SNAKE_CASE=7		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=True		,				SCREAMING_SNAKE_CASE=9_9		,				SCREAMING_SNAKE_CASE=3_2		,				SCREAMING_SNAKE_CASE=5		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=3_7		,				SCREAMING_SNAKE_CASE="gelu"		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=0.1		,				SCREAMING_SNAKE_CASE=5_1_2		,				SCREAMING_SNAKE_CASE=1_6		,				SCREAMING_SNAKE_CASE=2		,				SCREAMING_SNAKE_CASE=0.02		,				SCREAMING_SNAKE_CASE=3		,				SCREAMING_SNAKE_CASE=4		,				SCREAMING_SNAKE_CASE=None		,				) ->			Union[str, Any]:
  a__												= parent
  a__												= batch_size
  a__												= seq_length
  a__												= is_training
  a__												= use_token_type_ids
  a__												= use_labels
  a__												= vocab_size
  a__												= hidden_size
  a__												= num_hidden_layers
  a__												= num_attention_heads
  a__												= intermediate_size
  a__												= hidden_act
  a__												= hidden_dropout_prob
  a__												= attention_probs_dropout_prob
  a__												= max_position_embeddings
  a__												= type_vocab_size
  a__												= type_sequence_label_size
  a__												= initializer_range
  a__												= num_labels
  a__												= num_choices
  a__												= scope
  a__												= self.vocab_size - 1

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= ids_tensor([self.batch_size, self.seq_length]		,				self.vocab_size      )

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

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

  a__												= 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		,				)

  a__												= 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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			int:
  a__												= OpenAIGPTModel(config=SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				head_mask=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE      )
  a__												= model(SCREAMING_SNAKE_CASE      )

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

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Tuple:
  a__												= OpenAIGPTLMHeadModel(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Dict:
  a__												= OpenAIGPTDoubleHeadsModel(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

  a__												= model(SCREAMING_SNAKE_CASE		,				token_type_ids=SCREAMING_SNAKE_CASE		,				labels=SCREAMING_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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				*SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  a__												= self.num_labels
  a__												= OpenAIGPTForSequenceClassification(SCREAMING_SNAKE_CASE      )
  model.to(SCREAMING_SNAKE_CASE      )
  model.eval()

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



 def 						_UpperCAmelCase       (	self      ) ->			str:
  a__												= self.prepare_config_and_inputs()
  (
      (
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			(
      a__
  )							,			
  )												= config_and_inputs
  a__												= {
      '''input_ids''': input_ids,
      '''token_type_ids''': token_type_ids,
      '''head_mask''': head_mask,
  }

  return config, inputs_dict





@require_torch
class    __UpperCamelCase			(       _lowercase	,	_lowercase	,	_lowercase	,	unittest.TestCase				):

 """simple docstring"""



 _lowercase   :    Dict      						=					(
     (OpenAIGPTModel, OpenAIGPTLMHeadModel, OpenAIGPTDoubleHeadsModel, OpenAIGPTForSequenceClassification)
     if is_torch_available()
     else ()
 )
 _lowercase   :    Optional[Any]      						=					(
     (OpenAIGPTLMHeadModel,) if is_torch_available() else ()
 )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
 _lowercase   :    Dict      						=					(
     {
         '''feature-extraction''': OpenAIGPTModel,
         '''text-classification''': OpenAIGPTForSequenceClassification,
         '''text-generation''': OpenAIGPTLMHeadModel,
         '''zero-shot''': OpenAIGPTForSequenceClassification,
     }
     if is_torch_available()
     else {}
 )

 def 						_UpperCAmelCase       (	self		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE      ) ->			Optional[Any]:
  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		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE=False      ) ->			List[str]:
  a__												= super()._prepare_for_class(SCREAMING_SNAKE_CASE		,				SCREAMING_SNAKE_CASE		,				return_labels=SCREAMING_SNAKE_CASE      )

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

 def 						_UpperCAmelCase       (	self      ) ->			List[str]:
  a__												= OpenAIGPTModelTester(self      )
  a__												= ConfigTester(self		,				config_class=SCREAMING_SNAKE_CASE		,				n_embd=3_7      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[Any]:
  self.config_tester.run_common_tests()

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			int:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_lm_head_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_double_lm_head_model(*SCREAMING_SNAKE_CASE      )

 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*SCREAMING_SNAKE_CASE      )



 @slow
 def 						_UpperCAmelCase       (	self      ) ->			Union[str, Any]:
  for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
   a__												= OpenAIGPTModel.from_pretrained(SCREAMING_SNAKE_CASE      )
   self.assertIsNotNone(SCREAMING_SNAKE_CASE      )





@require_torch
class    __UpperCamelCase			(       unittest.TestCase				):

 """simple docstring"""



 @slow
 def 						_UpperCAmelCase       (	self      ) ->			Optional[int]:
  a__												= OpenAIGPTLMHeadModel.from_pretrained('''openai-gpt'''      )
  model.to(SCREAMING_SNAKE_CASE      )
  a__												= torch.tensor([[4_8_1, 4_7_3_5, 5_4_4]]		,				dtype=torch.long		,				device=SCREAMING_SNAKE_CASE      )  # the president is
  a__												= [
      4_8_1,
      4_7_3_5,
      5_4_4,
      2_4_6,
      9_6_3,
      8_7_0,
      7_6_2,
      2_3_9,
      2_4_4,
      4_0_4_7_7,
      2_4_4,
      2_4_9,
      7_1_9,
      8_8_1,
      4_8_7,
      5_4_4,
      2_4_0,
      2_4_4,
      6_0_3,
      4_8_1,
  ]  # the president is a very good man. " \n " i\'m sure he is, " said the

  a__												= model.generate(SCREAMING_SNAKE_CASE		,				do_sample=SCREAMING_SNAKE_CASE      )
  self.assertListEqual(output_ids[0].tolist()		,				SCREAMING_SNAKE_CASE      )





