




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




      """simple docstring"""





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

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    Optional[int]           =	ids_tensor([self.batch_size, self.seq_length]    ,				self.vocab_size     )

         snake_case				:    List[str]           =	None
         if self.use_token_type_ids:
            snake_case				:    Optional[Any]           =	ids_tensor([self.batch_size, self.seq_length]    ,				self.type_vocab_size     )

         snake_case				:    Optional[int]           =	None
         snake_case				:    Tuple           =	None
         snake_case				:    Optional[int]           =	None
         if self.use_labels:
            snake_case				:    Union[str, Any]           =	ids_tensor([self.batch_size]    ,				self.type_sequence_label_size     )
            snake_case				:    Optional[int]           =	ids_tensor([self.batch_size, self.seq_length]    ,				self.num_labels     )
            snake_case				:    Any           =	ids_tensor([self.batch_size]    ,				self.num_choices     )

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

         snake_case				:    str           =	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    ,				A    ,				A    ,				A    ,				A    ,				*A     )		->    Any:
         snake_case				:    Optional[Any]           =	OpenAIGPTModel(config=A     )
         model.to(A     )
         model.eval()

         snake_case				:    str           =	model(A    ,				token_type_ids=A    ,				head_mask=A     )
         snake_case				:    Any           =	model(A    ,				token_type_ids=A     )
         snake_case				:    Optional[int]           =	model(A     )

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

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				*A     )		->    List[Any]:
         snake_case				:    List[Any]           =	OpenAIGPTLMHeadModel(A     )
         model.to(A     )
         model.eval()

         snake_case				:    Any           =	model(A    ,				token_type_ids=A    ,				labels=A     )
         self.parent.assertEqual(result.loss.shape    ,				()     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				*A     )		->    List[Any]:
         snake_case				:    str           =	OpenAIGPTDoubleHeadsModel(A     )
         model.to(A     )
         model.eval()

         snake_case				:    List[str]           =	model(A    ,				token_type_ids=A    ,				labels=A     )
         self.parent.assertEqual(result.loss.shape    ,				()     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.seq_length, self.vocab_size)     )

      def      UpperCAmelCase				(				self    ,				A    ,				A    ,				A    ,				A    ,				*A     )		->    Union[str, Any]:
         snake_case				:    List[Any]           =	self.num_labels
         snake_case				:    Any           =	OpenAIGPTForSequenceClassification(A     )
         model.to(A     )
         model.eval()

         snake_case				:    Tuple           =	ids_tensor([self.batch_size]    ,				self.type_sequence_label_size     )
         snake_case				:    Optional[Any]           =	model(A    ,				token_type_ids=A    ,				labels=A     )
         self.parent.assertEqual(result.logits.shape    ,				(self.batch_size, self.num_labels)     )

      def      UpperCAmelCase				(				self     )		->    Union[str, Any]:
         snake_case				:    Union[str, Any]           =	self.prepare_config_and_inputs()
         (
             snake_case 
         )				:    List[str]           =	config_and_inputs
         snake_case				:    Tuple           =	{
             """input_ids""": input_ids,
             """token_type_ids""": token_type_ids,
             """head_mask""": head_mask,
         }

         return config, inputs_dict




@require_torch
class 						__lowercase		(UpperCamelCase__					,		UpperCamelCase__					,		UpperCamelCase__					,		unittest.TestCase							):




      """simple docstring"""





      _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    ,				A    ,				A    ,				A    ,				A    ,				A     )		->    Optional[int]:
         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    ,				A    ,				A    ,				A=False     )		->    int:
         snake_case				:    Union[str, Any]           =	super()._prepare_for_class(A    ,				A    ,				return_labels=A     )

         if return_labels:
            if model_class.__name__ == "OpenAIGPTDoubleHeadsModel":
               snake_case				:    Optional[int]           =	torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.num_choices, self.model_tester.seq_length)    ,				dtype=torch.long    ,				device=A    ,				)
               snake_case				:    List[str]           =	inputs_dict["""labels"""]
               snake_case				:    List[Any]           =	inputs_dict["""labels"""]
               snake_case				:    Tuple           =	torch.zeros(
                   (self.model_tester.batch_size, self.model_tester.num_choices)    ,				dtype=torch.long    ,				device=A    ,				)
               snake_case				:    List[str]           =	torch.zeros(
                   self.model_tester.batch_size    ,				dtype=torch.long    ,				device=A     )
         return inputs_dict

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    Any           =	OpenAIGPTModelTester(self     )
         snake_case				:    Tuple           =	ConfigTester(self    ,				config_class=A    ,				n_embd=3_7     )

      def      UpperCAmelCase				(				self     )		->    Any:
         self.config_tester.run_common_tests()

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    List[Any]           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_model(*A     )

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    Dict           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_lm_head_model(*A     )

      def      UpperCAmelCase				(				self     )		->    int:
         snake_case				:    List[str]           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_double_lm_head_model(*A     )

      def      UpperCAmelCase				(				self     )		->    Tuple:
         snake_case				:    str           =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*A     )

      @slow
      def      UpperCAmelCase				(				self     )		->    int:
         for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            snake_case				:    Optional[Any]           =	OpenAIGPTModel.from_pretrained(A     )
            self.assertIsNotNone(A     )




@require_torch
class 						__lowercase		(unittest.TestCase							):




      """simple docstring"""





      @slow
      def      UpperCAmelCase				(				self     )		->    List[Any]:
         snake_case				:    Dict           =	OpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""     )
         model.to(A     )
         snake_case				:    int           =	torch.tensor([[4_8_1, 4_7_3_5, 5_4_4]]    ,				dtype=torch.long    ,				device=A     )  # the president is
         snake_case				:    Optional[int]           =	[
             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

         snake_case				:    Tuple           =	model.generate(A    ,				do_sample=A     )
         self.assertListEqual(output_ids[0].tolist()    ,				A     )











from .glue import GlueDataset, GlueDataTrainingArguments
from .language_modeling import (
    LineByLineTextDataset,
    LineByLineWithRefDataset,
    LineByLineWithSOPTextDataset,
    TextDataset,
    TextDatasetForNextSentencePrediction,
)
from .squad import SquadDataset, SquadDataTrainingArguments





