



'''simple docstring'''







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






			def __init__(   self  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       ):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      parent
					SCREAMING_SNAKE_CASE_						:List[str]										=      13
					SCREAMING_SNAKE_CASE_						:Any										=      7
					SCREAMING_SNAKE_CASE_						:str										=      True
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      True
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      True
					SCREAMING_SNAKE_CASE_						:Dict										=      True
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      True
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      99
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      32
					SCREAMING_SNAKE_CASE_						:Tuple										=      2
					SCREAMING_SNAKE_CASE_						:List[str]										=      4
					SCREAMING_SNAKE_CASE_						:str										=      37
					SCREAMING_SNAKE_CASE_						:List[Any]										=      '''gelu'''
					SCREAMING_SNAKE_CASE_						:int										=      0.1
					SCREAMING_SNAKE_CASE_						:str										=      0.1
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      512
					SCREAMING_SNAKE_CASE_						:List[str]										=      16
					SCREAMING_SNAKE_CASE_						:Dict										=      2
					SCREAMING_SNAKE_CASE_						:Tuple										=      0.02
					SCREAMING_SNAKE_CASE_						:List[str]										=      3
					SCREAMING_SNAKE_CASE_						:int										=      4
					SCREAMING_SNAKE_CASE_						:Optional[int]										=      None
					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      self.vocab_size - 1






			def 				_lowercase	(   self  :							List[Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size	)

					SCREAMING_SNAKE_CASE_						:List[str]										=      None
					if self.use_input_mask:
							SCREAMING_SNAKE_CASE_						:List[str]										=      random_attention_mask([self.batch_size, self.seq_length]	)

					SCREAMING_SNAKE_CASE_						:Optional[int]										=      None
					if self.use_token_type_ids:
							SCREAMING_SNAKE_CASE_						:str										=      ids_tensor([self.batch_size, self.seq_length]      ,       self.type_vocab_size	)

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

					SCREAMING_SNAKE_CASE_						:int										=      None
					SCREAMING_SNAKE_CASE_						:List[str]										=      None
					SCREAMING_SNAKE_CASE_						:Dict										=      None
					if self.use_labels:
							SCREAMING_SNAKE_CASE_						:Any										=      ids_tensor([self.batch_size]      ,       self.type_sequence_label_size	)
							SCREAMING_SNAKE_CASE_						:str										=      ids_tensor([self.batch_size, self.seq_length]      ,       self.num_labels	)
							SCREAMING_SNAKE_CASE_						:Any										=      ids_tensor([self.batch_size]      ,       self.num_choices	)

					SCREAMING_SNAKE_CASE_						:Tuple										=      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_						:List[str]										=      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 				_lowercase	(   self  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							str      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							int      ,       SCREAMING_SNAKE_CASE  :							Optional[int]      ,       *SCREAMING_SNAKE_CASE  :							List[Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      TFOpenAIGPTModel(config=UpperCamelCase_	)
					SCREAMING_SNAKE_CASE_						:int										=      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					SCREAMING_SNAKE_CASE_						:Dict										=      model(UpperCamelCase_	)

					SCREAMING_SNAKE_CASE_						:str										=      [input_ids, input_mask]
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      model(UpperCamelCase_	)

					SCREAMING_SNAKE_CASE_						:List[Any]										=      model(UpperCamelCase_	)

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






			def 				_lowercase	(   self  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							str      ,       SCREAMING_SNAKE_CASE  :							Dict      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       *SCREAMING_SNAKE_CASE  :							List[str]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:List[str]										=      TFOpenAIGPTLMHeadModel(config=UpperCamelCase_	)
					SCREAMING_SNAKE_CASE_						:List[Any]										=      {'''input_ids''': input_ids, '''attention_mask''': input_mask, '''token_type_ids''': token_type_ids}
					SCREAMING_SNAKE_CASE_						:Tuple										=      model(UpperCamelCase_	)
					self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)	)






			def 				_lowercase	(   self  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							Dict      ,       SCREAMING_SNAKE_CASE  :							List[Any]      ,       SCREAMING_SNAKE_CASE  :							str      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Optional[int]      ,       SCREAMING_SNAKE_CASE  :							Dict      ,       *SCREAMING_SNAKE_CASE  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      TFOpenAIGPTDoubleHeadsModel(config=UpperCamelCase_	)

					SCREAMING_SNAKE_CASE_						:Tuple										=      tf.tile(tf.expand_dims(UpperCamelCase_      ,       1	)      ,       (1, self.num_choices, 1)	)
					SCREAMING_SNAKE_CASE_						:Tuple										=      tf.tile(tf.expand_dims(UpperCamelCase_      ,       1	)      ,       (1, self.num_choices, 1)	)
					SCREAMING_SNAKE_CASE_						:str										=      tf.tile(tf.expand_dims(UpperCamelCase_      ,       1	)      ,       (1, self.num_choices, 1)	)

					SCREAMING_SNAKE_CASE_						:Optional[int]										=      {
					    '''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,
					}
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      model(UpperCamelCase_	)
					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 				_lowercase	(   self  :							Dict      ,       SCREAMING_SNAKE_CASE  :							List[str]      ,       SCREAMING_SNAKE_CASE  :							str      ,       SCREAMING_SNAKE_CASE  :							Tuple      ,       SCREAMING_SNAKE_CASE  :							Dict      ,       SCREAMING_SNAKE_CASE  :							str      ,       *SCREAMING_SNAKE_CASE  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      self.num_labels
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      ids_tensor([self.batch_size]      ,       self.type_sequence_label_size	)
					SCREAMING_SNAKE_CASE_						:Dict										=      {
					    '''input_ids''': input_ids,
					    '''attention_mask''': input_mask,
					    '''token_type_ids''': token_type_ids,
					    '''labels''': sequence_labels,
					}
					SCREAMING_SNAKE_CASE_						:List[Any]										=      TFOpenAIGPTForSequenceClassification(UpperCamelCase_	)
					SCREAMING_SNAKE_CASE_						:Optional[Any]										=      model(UpperCamelCase_	)
					self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)	)






			def 				_lowercase	(   self  :							Union[str, Any]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Tuple										=      self.prepare_config_and_inputs()

					(
					    SCREAMING_SNAKE_CASE_   
					)						:Dict										=      config_and_inputs

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




@require_tf
class 				__lowerCAmelCase(							a__		,			a__		,			unittest.TestCase      ):
			__snake_case  :		List[Any]		=      (
			    (TFOpenAIGPTModel, TFOpenAIGPTLMHeadModel, TFOpenAIGPTDoubleHeadsModel, TFOpenAIGPTForSequenceClassification)
			    if is_tf_available()
			    else ()
			)
			__snake_case  :		List[Any]		=      (
			    (TFOpenAIGPTLMHeadModel,) if is_tf_available() else ()
			)  # TODO (PVP): Add Double HeadsModel when generate() function is changed accordingly
			__snake_case  :		Optional[int]		=      (
			    {
			        'feature-extraction': TFOpenAIGPTModel,
			        'text-classification': TFOpenAIGPTForSequenceClassification,
			        'text-generation': TFOpenAIGPTLMHeadModel,
			        'zero-shot': TFOpenAIGPTForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			__snake_case  :		Optional[Any]		=      False
			__snake_case  :		int		=      False






			def 				_lowercase	(   self  :							Union[str, Any]      ,       SCREAMING_SNAKE_CASE  :							Any      ,       SCREAMING_SNAKE_CASE  :							Optional[Any]      ,       SCREAMING_SNAKE_CASE  :							List[Any]      ,       SCREAMING_SNAKE_CASE  :							Union[str, Any]      ,       SCREAMING_SNAKE_CASE  :							List[str]	):




					"""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 				_lowercase	(   self  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:List[str]										=      TFOpenAIGPTModelTester(self	)
					SCREAMING_SNAKE_CASE_						:Any										=      ConfigTester(self      ,       config_class=UpperCamelCase_      ,       n_embd=37	)






			def 				_lowercase	(   self  :							Optional[Any]	):




					"""simple docstring"""







					self.config_tester.run_common_tests()






			def 				_lowercase	(   self  :							Optional[int]	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_openai_gpt_model(*UpperCamelCase_	)






			def 				_lowercase	(   self  :							str	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[int]										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_openai_gpt_lm_head(*UpperCamelCase_	)






			def 				_lowercase	(   self  :							str	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Any										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_openai_gpt_double_head(*UpperCamelCase_	)






			def 				_lowercase	(   self  :							Tuple	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:Optional[int]										=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*UpperCamelCase_	)






			@slow
			def 				_lowercase	(   self  :							Union[str, Any]	):




					"""simple docstring"""







					for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							SCREAMING_SNAKE_CASE_						:Union[str, Any]										=      TFOpenAIGPTModel.from_pretrained(UpperCamelCase_	)
							self.assertIsNotNone(UpperCamelCase_	)




@require_tf
class 				__lowerCAmelCase(							unittest.TestCase      ):






			@slow
			def 				_lowercase	(   self  :							int	):




					"""simple docstring"""







					SCREAMING_SNAKE_CASE_						:int										=      TFOpenAIGPTLMHeadModel.from_pretrained('openai-gpt'	)
					SCREAMING_SNAKE_CASE_						:Dict										=      tf.convert_to_tensor([[481, 4_735, 544]]      ,       dtype=tf.intaa	)  # the president is
					SCREAMING_SNAKE_CASE_						:List[str]										=      [
					    481,
					    4_735,
					    544,
					    246,
					    963,
					    870,
					    762,
					    239,
					    244,
					    40_477,
					    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_						:Tuple										=      model.generate(UpperCamelCase_      ,       do_sample=UpperCamelCase_	)
					self.assertListEqual(output_ids[0].numpy().tolist()      ,       UpperCamelCase_	)




"""simple docstring"""





import argparse
import time

import torch
import transformers
from measures_util import end_measure, log_measures, start_measure
from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForSeqaSeqLM, AutoTokenizer

from accelerate.utils import compute_module_sizes


snake_case__		:  Optional[Any]     					=					{
    '''gpt-j-6b''': {'''is_causal''': True, '''model''': '''sgugger/sharded-gpt-j-6B''', '''tokenizer''': '''EleutherAI/gpt-j-6B'''},
    '''gpt-neox''': {'''is_causal''': True, '''model''': '''EleutherAI/gpt-neox-20b'''},
    '''opt''': {'''is_causal''': True, '''model''': '''facebook/opt-30b'''},
    '''T0pp''': {'''is_causal''': False, '''model''': '''bigscience/T0pp''', '''model_revision''': '''sharded'''},
}

snake_case__		:  Any     					=					[
    '''Hello, my name is''',
    '''Are unicorns real? Unicorns are''',
    '''For the first time in several years,''',
    '''My name is Julien and I am''',
    '''The goal of life is''',
    '''Whenever I\'m sad, I like to''',
]


def   _snake_case    (    ):
      lowerCAmelCase       :						int							= argparse.ArgumentParser(description='''Run and time generations on a big model using Accelerate.'''      )
      parser.add_argument('''model_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the model to try.'''      )
      parser.add_argument(
          '''--tokenizer_name'''			, type=_snake_case			, default=_snake_case			, help='''The name of the tokenizer (if different from the model.'''      )
      parser.add_argument('''--is_causal'''			, type=_snake_case			, default=_snake_case			, help='''Whether or not the model is causal.'''      )
      parser.add_argument(
          '''--model_revision'''			, type=_snake_case			, default=_snake_case			, help='''The revision to use for the model checkpoint.'''      )
      parser.add_argument('''--torch_dtype'''			, type=_snake_case			, default=_snake_case			, help='''The dtype for the model.'''      )
      parser.add_argument('''--disk_offload'''			, action='''store_true'''      )

      lowerCAmelCase       :						Dict							= parser.parse_args()

      # Sanitize args
      if args.model_name in DEFAULT_MODELS:
            lowerCAmelCase       :						Optional[int]							= DEFAULT_MODELS[args.model_name]
            lowerCAmelCase       :						Union[str, Any]							= defaults['''model''']
            if args.tokenizer_name is None:
                  lowerCAmelCase       :						Union[str, Any]							= defaults.get('''tokenizer'''			, args.model_name      )
            if args.is_causal is None:
                  lowerCAmelCase       :						List[Any]							= defaults['''is_causal''']
            if args.model_revision is None:
                  lowerCAmelCase       :						Tuple							= defaults.get('''model_revision'''			, '''main'''      )

      if args.is_causal is None:
            raise ValueError('''Could not infer the default for `--is_causal`, pass either True or False for it.'''      )
      if args.tokenizer_name is None:
            lowerCAmelCase       :						Union[str, Any]							= args.model_name
      if args.model_revision is None:
            lowerCAmelCase       :						List[Any]							= '''main'''

      return args


def   _snake_case    (    ):
      transformers.utils.logging.set_verbosity_error()
      lowerCAmelCase       :						Union[str, Any]							= parse_args()

      if args.torch_dtype is None:
            lowerCAmelCase       :						List[str]							= AutoConfig.from_pretrained(args.model_name      )
            lowerCAmelCase       :						Optional[Any]							= getattr(_snake_case			, '''torch_dtype'''			, torch.floataa      )
      else:
            lowerCAmelCase       :						Optional[int]							= getattr(_snake_case			, args.torch_dtype      )
      lowerCAmelCase       :						Tuple							= AutoModelForCausalLM if args.is_causal else AutoModelForSeqaSeqLM
      lowerCAmelCase       :						str							= {
          '''torch_dtype''': torch_dtype,
          '''revision''': args.model_revision,
      }
      if args.disk_offload:
            lowerCAmelCase       :						Optional[int]							= '''tmp_offload'''
            lowerCAmelCase       :						Optional[int]							= True

      lowerCAmelCase       :						str							= start_measure()
      lowerCAmelCase       :						List[str]							= model_cls.from_pretrained(args.model_name			, device_map='''auto'''			, **_snake_case      )
      lowerCAmelCase       :						Tuple							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model loading'''      )

      lowerCAmelCase       :						str							= compute_module_sizes(_snake_case      )
      lowerCAmelCase       :						Tuple							= {v: 0 for v in model.hf_device_map.values()}
      for module, device in model.hf_device_map.items():
            device_size[device] += module_sizes[module]
      lowerCAmelCase       :						int							= '''\n'''.join([f'''- {device}: {size // 2**20}MiB''' for device, size in device_size.items()]      )
      print(f'''\nTheoretical use:\n{message}'''      )

      lowerCAmelCase       :						int							= AutoTokenizer.from_pretrained(args.tokenizer_name      )

      lowerCAmelCase       :						List[str]							= start_measure()
      lowerCAmelCase       :						Union[str, Any]							= []
      lowerCAmelCase       :						int							= []
      lowerCAmelCase       :						Dict							= []
      for prompt in PROMPTS:
            lowerCAmelCase       :						str							= tokenizer(_snake_case			, return_tensors='''pt'''      ).to(0      )
            lowerCAmelCase       :						int							= inputs['''input_ids'''][0].tolist()
            lowerCAmelCase       :						Optional[Any]							= time.time()
            lowerCAmelCase       :						Optional[int]							= model.generate(inputs['''input_ids''']      )
            lowerCAmelCase       :						Any							= time.time()
            lowerCAmelCase       :						str							= outputs[0].tolist()
            lowerCAmelCase       :						int							= len(_snake_case      ) if outputs[: len(_snake_case      )] != tokens else len(_snake_case      ) - len(_snake_case      )
            lowerCAmelCase       :						Tuple							= after_generate - before_generate

            lowerCAmelCase       :						Tuple							= tokenizer.decode(_snake_case			, skip_special_tokens=_snake_case      )
            texts_outs.append(_snake_case      )
            generation_times.append(_snake_case      )
            gen_tokens.append(_snake_case      )
            print(f'''Prompt: {prompt}\nGeneration {text_out}\nIn {generation_time:.2f}s for {num_gen_tokens} tokens\n'''      )

      lowerCAmelCase       :						List[Any]							= end_measure(_snake_case      )
      log_measures(_snake_case			, '''Model generation'''      )

      lowerCAmelCase       :						str							= [gen / tok for gen, tok in zip(_snake_case			, _snake_case      )]
      lowerCAmelCase       :						Any							= sum(_snake_case      ) / len(_snake_case      )
      print(f'''Average time of generation per token: {avg_gen:.2f}s'''      )
      print(f'''First generation (avg time per token): {generation_times_per_token[0]:.2f}s'''      )
      lowerCAmelCase       :						List[Any]							= sum(generation_times_per_token[1:]      ) / (len(_snake_case      ) - 1)
      print(f'''Average time of generation per token (excluding the first): {avg_gen:.2f}s'''      )


if __name__ == "__main__":
 main()


