

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







   '''simple docstring'''




   def __init__(       self					:			Optional[Any]     , __lowerCamelCase					:			Tuple     , __lowerCamelCase					:			Tuple=13     , __lowerCamelCase					:			int=7     , __lowerCamelCase					:			Any=True     , __lowerCamelCase					:			List[str]=True     , __lowerCamelCase					:			Tuple=True     , __lowerCamelCase					:			int=99     , __lowerCamelCase					:			Optional[Any]=32     , __lowerCamelCase					:			str=5     , __lowerCamelCase					:			Optional[Any]=4     , __lowerCamelCase					:			Any=37     , __lowerCamelCase					:			Optional[int]="gelu"     , __lowerCamelCase					:			List[str]=0.1     , __lowerCamelCase					:			Any=0.1     , __lowerCamelCase					:			Tuple=512     , __lowerCamelCase					:			Dict=16     , __lowerCamelCase					:			str=2     , __lowerCamelCase					:			List[str]=0.02     , __lowerCamelCase					:			List[Any]=3     , __lowerCamelCase					:			int=4     , __lowerCamelCase					:			List[Any]=None     , ):
       SCREAMING_SNAKE_CASE										=   parent
       SCREAMING_SNAKE_CASE										=   batch_size
       SCREAMING_SNAKE_CASE										=   seq_length
       SCREAMING_SNAKE_CASE										=   is_training
       SCREAMING_SNAKE_CASE										=   use_token_type_ids
       SCREAMING_SNAKE_CASE										=   use_labels
       SCREAMING_SNAKE_CASE										=   vocab_size
       SCREAMING_SNAKE_CASE										=   hidden_size
       SCREAMING_SNAKE_CASE										=   num_hidden_layers
       SCREAMING_SNAKE_CASE										=   num_attention_heads
       SCREAMING_SNAKE_CASE										=   intermediate_size
       SCREAMING_SNAKE_CASE										=   hidden_act
       SCREAMING_SNAKE_CASE										=   hidden_dropout_prob
       SCREAMING_SNAKE_CASE										=   attention_probs_dropout_prob
       SCREAMING_SNAKE_CASE										=   max_position_embeddings
       SCREAMING_SNAKE_CASE										=   type_vocab_size
       SCREAMING_SNAKE_CASE										=   type_sequence_label_size
       SCREAMING_SNAKE_CASE										=   initializer_range
       SCREAMING_SNAKE_CASE										=   num_labels
       SCREAMING_SNAKE_CASE										=   num_choices
       SCREAMING_SNAKE_CASE										=   scope
       SCREAMING_SNAKE_CASE										=   self.vocab_size - 1





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   ids_tensor([self.batch_size, self.seq_length]     , self.vocab_size			)

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

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

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

       SCREAMING_SNAKE_CASE										=   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      _snake_case		(       self					:			Tuple     , __lowerCamelCase					:			Union[str, Any]     , __lowerCamelCase					:			List[str]     , __lowerCamelCase					:			Optional[int]     , __lowerCamelCase					:			Dict     , *__lowerCamelCase					:			Optional[Any]			):
       SCREAMING_SNAKE_CASE										=   OpenAIGPTModel(config=__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

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

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





   def      _snake_case		(       self					:			Dict     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			Union[str, Any]     , __lowerCamelCase					:			Tuple     , __lowerCamelCase					:			Dict     , *__lowerCamelCase					:			Tuple			):
       SCREAMING_SNAKE_CASE										=   OpenAIGPTLMHeadModel(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

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





   def      _snake_case		(       self					:			str     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			List[str]     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			Tuple     , *__lowerCamelCase					:			Tuple			):
       SCREAMING_SNAKE_CASE										=   OpenAIGPTDoubleHeadsModel(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

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





   def      _snake_case		(       self					:			Dict     , __lowerCamelCase					:			Optional[Any]     , __lowerCamelCase					:			Tuple     , __lowerCamelCase					:			Any     , __lowerCamelCase					:			Union[str, Any]     , *__lowerCamelCase					:			str			):
       SCREAMING_SNAKE_CASE										=   self.num_labels
       SCREAMING_SNAKE_CASE										=   OpenAIGPTForSequenceClassification(__UpperCamelCase			)
       model.to(__UpperCamelCase			)
       model.eval()

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





   def      _snake_case		(       self					:			Dict			):
       SCREAMING_SNAKE_CASE										=   self.prepare_config_and_inputs()
       (
           (
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	(
           SCREAMING_SNAKE_CASE   
       )							,	   
       )										=   config_and_inputs
       SCREAMING_SNAKE_CASE										=   {
           "input_ids": input_ids,
           "token_type_ids": token_type_ids,
           "head_mask": head_mask,
       }

       return config, inputs_dict





@require_torch
class 			_SCREAMING_SNAKE_CASE (   _lowerCamelCase					,       _lowerCamelCase					,       _lowerCamelCase					,       unittest.TestCase		):







   '''simple docstring'''




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





   def      _snake_case		(       self					:			Optional[int]     , __lowerCamelCase					:			Any     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			Dict     , __lowerCamelCase					:			str     , __lowerCamelCase					:			List[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      _snake_case		(       self					:			List[Any]     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			List[Any]     , __lowerCamelCase					:			Optional[Any]=False			):
       SCREAMING_SNAKE_CASE										=   super()._prepare_for_class(__UpperCamelCase     , __UpperCamelCase     , return_labels=__UpperCamelCase			)

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





   def      _snake_case		(       self					:			Any			):
       SCREAMING_SNAKE_CASE										=   OpenAIGPTModelTester(self			)
       SCREAMING_SNAKE_CASE										=   ConfigTester(self     , config_class=__UpperCamelCase     , n_embd=37			)





   def      _snake_case		(       self					:			str			):
       self.config_tester.run_common_tests()





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_openai_gpt_model(*__UpperCamelCase			)





   def      _snake_case		(       self					:			Union[str, Any]			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_lm_head_model(*__UpperCamelCase			)





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_double_lm_head_model(*__UpperCamelCase			)





   def      _snake_case		(       self					:			List[str]			):
       SCREAMING_SNAKE_CASE										=   self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__UpperCamelCase			)





   @slow
   def      _snake_case		(       self					:			Dict			):
       for model_name in OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
           SCREAMING_SNAKE_CASE										=   OpenAIGPTModel.from_pretrained(__UpperCamelCase			)
           self.assertIsNotNone(__UpperCamelCase			)





@require_torch
class 			_SCREAMING_SNAKE_CASE (   unittest.TestCase		):







   '''simple docstring'''




   @slow
   def      _snake_case		(       self					:			Optional[int]			):
       SCREAMING_SNAKE_CASE										=   OpenAIGPTLMHeadModel.from_pretrained("openai-gpt"			)
       model.to(__UpperCamelCase			)
       SCREAMING_SNAKE_CASE										=   torch.tensor([[481, 4735, 544]]     , dtype=torch.long     , device=__UpperCamelCase			)  # the president is
       SCREAMING_SNAKE_CASE										=   [
           481,
           4735,
           544,
           246,
           963,
           870,
           762,
           239,
           244,
           40477,
           244,
           249,
           719,
           881,
           487,
           544,
           240,
           244,
           603,
           481,
       ]  # the president is a very good man. " \n " i\'m sure he is, " said the

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




from typing import Callable, Optional

from .. import Features
from ..packaged_modules.generator.generator import Generator
from .abc import AbstractDatasetInputStream



class        lowerCAmelCase__					(  _lowerCamelCase       ):


				def __init__(							self      :  int						, __UpperCamelCase      :  Callable						, __UpperCamelCase      :  Optional[Features] = None						, __UpperCamelCase      :  str = None						, __UpperCamelCase      :  bool = False						, __UpperCamelCase      :  bool = False						, __UpperCamelCase      :  Optional[dict] = None						, __UpperCamelCase      :  Optional[int] = None						, **__UpperCamelCase      :  str						, )		->     List[str]:
											super().__init__(
											    features=__UpperCamelCase						, cache_dir=__UpperCamelCase						, keep_in_memory=__UpperCamelCase						, streaming=__UpperCamelCase						, num_proc=__UpperCamelCase						, **__UpperCamelCase						, )
											A      						=    Generator(
											    cache_dir=__UpperCamelCase						, features=__UpperCamelCase						, generator=__UpperCamelCase						, gen_kwargs=__UpperCamelCase						, **__UpperCamelCase						, )


				def __UpperCamelCase						(							self      :  Union[str, Any]       )		->     Tuple:
											# Build iterable dataset
											if self.streaming:
																		A      						=    self.builder.as_streaming_dataset(split='train'       )
											# Build regular (map-style) dataset
											else:
																		A      						=    None
																		A      						=    None
																		A      						=    None
																		A      						=    None

																		self.builder.download_and_prepare(
																		    download_config=__UpperCamelCase						, download_mode=__UpperCamelCase						, verification_mode=__UpperCamelCase						, base_path=__UpperCamelCase						, num_proc=self.num_proc						, )
																		A      						=    self.builder.as_dataset(
																		    split='train'						, verification_mode=__UpperCamelCase						, in_memory=self.keep_in_memory       )
											return dataset