


from __future__ import annotations

import unittest

from transformers import OpenAIGPTConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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


if is_tf_available():
   import tensorflow as tf

   from transformers.models.openai.modeling_tf_openai import (
       TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
       TFOpenAIGPTDoubleHeadsModel,
       TFOpenAIGPTForSequenceClassification,
       TFOpenAIGPTLMHeadModel,
       TFOpenAIGPTModel,
   )






class        UpperCAmelCase__			:



    """simple docstring"""






    def __init__(     self   ,  A_   ,  )    ->			int:
         __UpperCamelCase       				=parent
         __UpperCamelCase       				=13
         __UpperCamelCase       				=7
         __UpperCamelCase       				=True
         __UpperCamelCase       				=True
         __UpperCamelCase       				=True
         __UpperCamelCase       				=True
         __UpperCamelCase       				=True
         __UpperCamelCase       				=99
         __UpperCamelCase       				=32
         __UpperCamelCase       				=2
         __UpperCamelCase       				=4
         __UpperCamelCase       				=37
         __UpperCamelCase       				='gelu'
         __UpperCamelCase       				=0.1
         __UpperCamelCase       				=0.1
         __UpperCamelCase       				=512
         __UpperCamelCase       				=16
         __UpperCamelCase       				=2
         __UpperCamelCase       				=0.02
         __UpperCamelCase       				=3
         __UpperCamelCase       				=4
         __UpperCamelCase       				=None
         __UpperCamelCase       				=self.vocab_size - 1






    def _a   (     self       )    ->			Tuple:
         __UpperCamelCase       				=ids_tensor([self.batch_size, self.seq_length]   ,  self.vocab_size       )

         __UpperCamelCase       				=None
         if self.use_input_mask:
              __UpperCamelCase       				=random_attention_mask([self.batch_size, self.seq_length]       )

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

         __UpperCamelCase       				=None
         if self.use_mc_token_ids:
              __UpperCamelCase       				=ids_tensor([self.batch_size, self.num_choices]   ,  self.seq_length       )

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

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

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

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






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_   ,  *A_       )    ->			Dict:
         __UpperCamelCase       				=TFOpenAIGPTModel(config=A_       )
         __UpperCamelCase       				={'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
         __UpperCamelCase       				=model(A_       )

         __UpperCamelCase       				=[input_ids, input_mask]
         __UpperCamelCase       				=model(A_       )

         __UpperCamelCase       				=model(A_       )

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






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_   ,  *A_       )    ->			Optional[Any]:
         __UpperCamelCase       				=TFOpenAIGPTLMHeadModel(config=A_       )
         __UpperCamelCase       				={'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
         __UpperCamelCase       				=model(A_       )
         self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.seq_length, self.vocab_size)       )






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_   ,  *A_       )    ->			Optional[int]:
         __UpperCamelCase       				=TFOpenAIGPTDoubleHeadsModel(config=A_       )

         __UpperCamelCase       				=tf.tile(tf.expand_dims(A_   ,  1       )   ,  (1, self.num_choices, 1)       )
         __UpperCamelCase       				=tf.tile(tf.expand_dims(A_   ,  1       )   ,  (1, self.num_choices, 1)       )
         __UpperCamelCase       				=tf.tile(tf.expand_dims(A_   ,  1       )   ,  (1, self.num_choices, 1)       )

         __UpperCamelCase       				={
             'input_ids': multiple_choice_inputs_ids,
             'mc_token_ids': mc_token_ids,
             'attention_mask': multiple_choice_input_mask,
             'token_type_ids': multiple_choice_token_type_ids,
         }
         __UpperCamelCase       				=model(A_       )
         self.parent.assertEqual(
             result.logits.shape   ,  (self.batch_size, self.num_choices, self.seq_length, self.vocab_size)       )
         self.parent.assertEqual(result.mc_logits.shape   ,  (self.batch_size, self.num_choices)       )






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_   ,  *A_       )    ->			Union[str, Any]:
         __UpperCamelCase       				=self.num_labels
         __UpperCamelCase       				=ids_tensor([self.batch_size]   ,  self.type_sequence_label_size       )
         __UpperCamelCase       				={
             'input_ids': input_ids,
             'attention_mask': input_mask,
             'token_type_ids': token_type_ids,
             'labels': sequence_labels,
         }
         __UpperCamelCase       				=TFOpenAIGPTForSequenceClassification(A_       )
         __UpperCamelCase       				=model(A_       )
         self.parent.assertEqual(result.logits.shape   ,  (self.batch_size, self.num_labels)       )






    def _a   (     self       )    ->			List[Any]:
         __UpperCamelCase       				=self.prepare_config_and_inputs()

         (
             (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, (
             __UpperCamelCase
         )				, 
         )       				=config_and_inputs

         __UpperCamelCase       				={'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': input_mask}
         return config, inputs_dict






@require_tf
class        UpperCAmelCase__			(   A_  ,    A_  ,    unittest.TestCase						):



    """simple docstring"""






    UpperCAmelCase__		:						str    =    (
        (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
        if is_tf_available()
        else ()
    )
    UpperCAmelCase__		:						List[Any]    =    (
        (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
    )  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
    UpperCAmelCase__		:						Optional[int]    =    (
        {
            "feature-extraction": TFOpenAIGPTModel,
            "text-classification": TFOpenAIGPTForSequenceClassification,
            "text-generation": TFOpenAIGPTLMHeadModel,
            "zero-shot": TFOpenAIGPTForSequenceClassification,
        }
        if is_tf_available()
        else {}
    )
    UpperCAmelCase__		:						Any    =    False
    UpperCAmelCase__		:						Any    =    False






    def _a   (     self   ,  A_   ,  A_   ,  A_   ,  A_   ,  A_       )    ->			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 _a   (     self       )    ->			Dict:
         __UpperCamelCase       				=TFOpenAIGPTModelTester(self       )
         __UpperCamelCase       				=ConfigTester(self   ,  config_class=A_   ,  n_embd=37       )






    def _a   (     self       )    ->			Dict:
         self.config_tester.run_common_tests()






    def _a   (     self       )    ->			Optional[int]:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_model(*A_       )






    def _a   (     self       )    ->			List[str]:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_lm_head(*A_       )






    def _a   (     self       )    ->			Dict:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_double_head(*A_       )






    def _a   (     self       )    ->			Any:
         __UpperCamelCase       				=self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*A_       )






    @slow
    def _a   (     self       )    ->			Tuple:
         for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              __UpperCamelCase       				=TFOpenAIGPTModel.from_pretrained(A_       )
              self.assertIsNotNone(A_       )






@require_tf
class        UpperCAmelCase__			(   unittest.TestCase						):



    """simple docstring"""






    @slow
    def _a   (     self       )    ->			str:
         __UpperCamelCase       				=TFOpenAIGPTLMHeadModel.from_pretrained('openai-gpt'       )
         __UpperCamelCase       				=tf.convert_to_tensor([[481, 4735, 544]]   ,  dtype=tf.intaa       )  # the president is
         __UpperCamelCase       				=[
             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

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









import argparse
import logging

import numpy as np
import torch

from transformers import AutoModelForCausalLM, AutoTokenizer


logging.basicConfig(
    format='%(asctime)s - %(levelname)s - %(name)s - %(message)s',
    datefmt='%m/%d/%Y %H:%M:%S',
    level=logging.INFO,
)
_A									=		logging.getLogger(__name__)

def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   str							):
     np.random.seed(args.seed							)
     torch.manual_seed(args.seed							)
     if args.n_gpu > 0:
          torch.cuda.manual_seed_all(args.seed							)

def 		_UpperCAmelCase							(     ):
     __UpperCamelCase       				=argparse.ArgumentParser()
     parser.add_argument(
         '--model_name_or_path'     ,	default=SCREAMING_SNAKE_CASE__     ,	type=SCREAMING_SNAKE_CASE__     ,	required=SCREAMING_SNAKE_CASE__     ,	)
     parser.add_argument('--prompt'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=''							)
     parser.add_argument('--length'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=20							)
     parser.add_argument('--stop_token'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=SCREAMING_SNAKE_CASE__     ,	help='Token at which text generation is stopped'							)
     parser.add_argument(
         '--temperature'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=1.0     ,	help='temperature of 1.0 has no effect, lower tend toward greedy sampling'     ,	)
     parser.add_argument(
         '--repetition_penalty'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=1.0     ,	help='primarily useful for CTRL model; in that case, use 1.2'							)
     parser.add_argument('--k'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=0							)
     parser.add_argument('--penalty_alpha'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=0.0							)
     parser.add_argument('--p'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=0.9							)

     parser.add_argument('--prefix'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=''     ,	help='Text added prior to input.'							)
     parser.add_argument('--padding_text'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=''     ,	help='Deprecated, the use of `--prefix` is preferred.'							)
     parser.add_argument('--xlm_language'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=''     ,	help='Optional language when used with the XLM model.'							)

     parser.add_argument('--seed'     ,	type=SCREAMING_SNAKE_CASE__     ,	default=42     ,	help='random seed for initialization'							)
     parser.add_argument('--no_cuda'     ,	action='store_true'     ,	help='Avoid using CUDA when available'							)
     parser.add_argument(
         '--fp16'     ,	action='store_true'     ,	help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'     ,	)
     __UpperCamelCase       				=parser.parse_args()

     __UpperCamelCase       				=torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'							)
     __UpperCamelCase       				=0 if args.no_cuda else torch.cuda.device_count()

     logger.warning(F'device: {args.device}, n_gpu: {args.n_gpu}, 16-bits training: {args.fpaa}'							)

     set_seed(SCREAMING_SNAKE_CASE__							)

     # Initialize the model and tokenizer
     __UpperCamelCase       				=AutoTokenizer.from_pretrained(args.model_name_or_path							)
     __UpperCamelCase       				=AutoModelForCausalLM.from_pretrained(args.model_name_or_path							)

     # tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path)
     # model = OPTForCausalLM.from_pretrained(args.model_name_or_path)
     model.to(args.device							)

     if args.fpaa:
          model.half()

     logger.info(SCREAMING_SNAKE_CASE__							)
     __UpperCamelCase       				=args.prompt if args.prompt else input('Model prompt >>> '							)

     __UpperCamelCase       				=tokenizer(SCREAMING_SNAKE_CASE__     ,	return_tensors='pt'     ,	add_special_tokens=SCREAMING_SNAKE_CASE__							)
     __UpperCamelCase       				={key: value.to(args.device							) for key, value in inputs.items()}

     __UpperCamelCase       				=model.generate(
         **SCREAMING_SNAKE_CASE__     ,	max_length=args.length + len(inputs['input_ids'][0]							)     ,	penalty_alpha=args.penalty_alpha     ,	top_k=args.k     ,	)

     __UpperCamelCase       				=[]
     for generated_sequence_idx, generated_sequence in enumerate(SCREAMING_SNAKE_CASE__							):
          print(F'=== GENERATED SEQUENCE {generated_sequence_idx + 1} ==='							)
          __UpperCamelCase       				=generated_sequence.tolist()

          # Decode text
          __UpperCamelCase       				=tokenizer.decode(SCREAMING_SNAKE_CASE__     ,	clean_up_tokenization_spaces=SCREAMING_SNAKE_CASE__     ,	add_special_tokens=SCREAMING_SNAKE_CASE__							)

          # Remove all text after the stop token
          __UpperCamelCase       				=text[: text.find(args.stop_token							) if args.stop_token else None]

          # Add the prompt at the beginning of the sequence. Remove the excess text that was used for pre-processing
          __UpperCamelCase       				=(
              prompt_text + text[len(tokenizer.decode(inputs['input_ids'][0]     ,	clean_up_tokenization_spaces=SCREAMING_SNAKE_CASE__							)							) :]
          )

          generated_sequences.append(SCREAMING_SNAKE_CASE__							)
          print(SCREAMING_SNAKE_CASE__							)

     return generated_sequences


if __name__ == "__main__":
   main()





