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





							def __init__(					self    :  str	,					_UpperCAmelCase    :  Union[str, Any]	,					)		->						Optional[int]:



											"""simple docstring"""


											__lowercase								=   parent
											__lowercase								=   13
											__lowercase								=   7
											__lowercase								=   True
											__lowercase								=   True
											__lowercase								=   True
											__lowercase								=   True
											__lowercase								=   True
											__lowercase								=   99
											__lowercase								=   32
											__lowercase								=   2
											__lowercase								=   4
											__lowercase								=   37
											__lowercase								=   'gelu'
											__lowercase								=   0.1
											__lowercase								=   0.1
											__lowercase								=   5_12
											__lowercase								=   16
											__lowercase								=   2
											__lowercase								=   0.02
											__lowercase								=   3
											__lowercase								=   4
											__lowercase								=   None
											__lowercase								=   self.vocab_size - 1





							def      a__					(					self    :  int							)		->						List[str]:



											"""simple docstring"""


											__lowercase								=   ids_tensor([self.batch_size, self.seq_length]	,					self.vocab_size							)

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

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

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

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

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

											__lowercase								=   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    :  Dict	,					_UpperCAmelCase    :  Tuple	,					_UpperCAmelCase    :  Optional[int]	,					_UpperCAmelCase    :  Union[str, Any]	,					_UpperCAmelCase    :  str	,					_UpperCAmelCase    :  List[Any]	,					*_UpperCAmelCase    :  Any							)		->						Optional[Any]:



											"""simple docstring"""


											__lowercase								=   TFOpenAIGPTModel(config=_UpperCAmelCase							)
											__lowercase								=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
											__lowercase								=   model(_UpperCAmelCase							)

											__lowercase								=   [input_ids, input_mask]
											__lowercase								=   model(_UpperCAmelCase							)

											__lowercase								=   model(_UpperCAmelCase							)

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





							def      a__					(					self    :  Tuple	,					_UpperCAmelCase    :  Optional[Any]	,					_UpperCAmelCase    :  Optional[int]	,					_UpperCAmelCase    :  List[Any]	,					_UpperCAmelCase    :  Dict	,					_UpperCAmelCase    :  Optional[int]	,					*_UpperCAmelCase    :  Dict							)		->						Optional[Any]:



											"""simple docstring"""


											__lowercase								=   TFOpenAIGPTLMHeadModel(config=_UpperCAmelCase							)
											__lowercase								=   {'input_ids': input_ids, 'attention_mask': input_mask, 'token_type_ids': token_type_ids}
											__lowercase								=   model(_UpperCAmelCase							)
											self.parent.assertEqual(result.logits.shape	,					(self.batch_size, self.seq_length, self.vocab_size)							)





							def      a__					(					self    :  List[str]	,					_UpperCAmelCase    :  Any	,					_UpperCAmelCase    :  List[Any]	,					_UpperCAmelCase    :  Tuple	,					_UpperCAmelCase    :  Dict	,					_UpperCAmelCase    :  Optional[int]	,					_UpperCAmelCase    :  Tuple	,					*_UpperCAmelCase    :  List[str]							)		->						str:



											"""simple docstring"""


											__lowercase								=   TFOpenAIGPTDoubleHeadsModel(config=_UpperCAmelCase							)

											__lowercase								=   tf.tile(tf.expand_dims(_UpperCAmelCase	,					1							)	,					(1, self.num_choices, 1)							)
											__lowercase								=   tf.tile(tf.expand_dims(_UpperCAmelCase	,					1							)	,					(1, self.num_choices, 1)							)
											__lowercase								=   tf.tile(tf.expand_dims(_UpperCAmelCase	,					1							)	,					(1, self.num_choices, 1)							)

											__lowercase								=   {
											    '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,
											}
											__lowercase								=   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      a__					(					self    :  Tuple	,					_UpperCAmelCase    :  List[str]	,					_UpperCAmelCase    :  List[str]	,					_UpperCAmelCase    :  Optional[int]	,					_UpperCAmelCase    :  List[Any]	,					_UpperCAmelCase    :  List[Any]	,					*_UpperCAmelCase    :  Dict							)		->						List[Any]:



											"""simple docstring"""


											__lowercase								=   self.num_labels
											__lowercase								=   ids_tensor([self.batch_size]	,					self.type_sequence_label_size							)
											__lowercase								=   {
											    'input_ids': input_ids,
											    'attention_mask': input_mask,
											    'token_type_ids': token_type_ids,
											    'labels': sequence_labels,
											}
											__lowercase								=   TFOpenAIGPTForSequenceClassification(_UpperCAmelCase							)
											__lowercase								=   model(_UpperCAmelCase							)
											self.parent.assertEqual(result.logits.shape	,					(self.batch_size, self.num_labels)							)





							def      a__					(					self    :  Tuple							)		->						Tuple:



											"""simple docstring"""


											__lowercase								=   self.prepare_config_and_inputs()

											(
											    (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   (
											    __lowercase
											)						,   
											)								=   config_and_inputs

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






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





							def      a__					(					self    :  List[str]	,					_UpperCAmelCase    :  Union[str, Any]	,					_UpperCAmelCase    :  Any	,					_UpperCAmelCase    :  int	,					_UpperCAmelCase    :  str	,					_UpperCAmelCase    :  int							)		->						Union[str, Any]:



											"""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      a__					(					self    :  Optional[int]							)		->						Any:



											"""simple docstring"""


											__lowercase								=   TFOpenAIGPTModelTester(self							)
											__lowercase								=   ConfigTester(self	,					config_class=_UpperCAmelCase	,					n_embd=37							)





							def      a__					(					self    :  Optional[Any]							)		->						List[Any]:



											"""simple docstring"""


											self.config_tester.run_common_tests()





							def      a__					(					self    :  int							)		->						str:



											"""simple docstring"""


											__lowercase								=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_openai_gpt_model(*_UpperCAmelCase							)





							def      a__					(					self    :  Optional[int]							)		->						Dict:



											"""simple docstring"""


											__lowercase								=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_openai_gpt_lm_head(*_UpperCAmelCase							)





							def      a__					(					self    :  int							)		->						Union[str, Any]:



											"""simple docstring"""


											__lowercase								=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_openai_gpt_double_head(*_UpperCAmelCase							)





							def      a__					(					self    :  Union[str, Any]							)		->						Dict:



											"""simple docstring"""


											__lowercase								=   self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*_UpperCAmelCase							)





							@slow
							def      a__					(					self    :  List[Any]							)		->						Union[str, Any]:



											"""simple docstring"""


											for model_name in TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__lowercase								=   TFOpenAIGPTModel.from_pretrained(_UpperCAmelCase							)
															self.assertIsNotNone(_UpperCAmelCase							)






@require_tf
class       A__  (		unittest.TestCase							):





							@slow
							def      a__					(					self    :  str							)		->						Tuple:



											"""simple docstring"""


											__lowercase								=   TFOpenAIGPTLMHeadModel.from_pretrained('openai-gpt'							)
											__lowercase								=   tf.convert_to_tensor([[4_81, 47_35, 5_44]]	,					dtype=tf.intaa							)  # the president is
											__lowercase								=   [
											    4_81,
											    47_35,
											    5_44,
											    2_46,
											    9_63,
											    8_70,
											    7_62,
											    2_39,
											    2_44,
											    4_04_77,
											    2_44,
											    2_49,
											    7_19,
											    8_81,
											    4_87,
											    5_44,
											    2_40,
											    2_44,
											    6_03,
											    4_81,
											]  # the president is a very good man. " \n " i\'m sure he is, " said the

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







import logging
from pathlib import Path

import numpy as np
import pytorch_lightning as pl
import torch
from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
from pytorch_lightning.utilities import rank_zero_only

from utils import save_json




def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: List[Any]     )	->		str:
				__lowercase								=   filter(lambda SCREAMING_SNAKE_CASE     : p.requires_grad       ,				model.parameters()     )
				__lowercase								=   sum([np.prod(p.size()     ) for p in model_parameters]     )
				return params


SCREAMING_SNAKE_CASE__				=							logging.getLogger(__name__)
class       A__  (		pl.Callback							):





							def      a__					(					self    :  List[str]	,					_UpperCAmelCase    :  Tuple	,					_UpperCAmelCase    :  List[str]							)		->						Tuple:



											"""simple docstring"""


											__lowercase								=   {f"""lr_group_{i}""": param['lr'] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups							)}
											pl_module.logger.log_metrics(_UpperCAmelCase							)





							@rank_zero_only
							def      a__					(					self    :  Dict	,					_UpperCAmelCase    :  pl.Trainer	,					_UpperCAmelCase    :  pl.LightningModule	,					_UpperCAmelCase    :  str	,					_UpperCAmelCase    :  int=True							)		->						None:



											"""simple docstring"""


											logger.info(f"""***** {type_path} results at step {trainer.global_step:05d} *****"""							)
											__lowercase								=   trainer.callback_metrics
											trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ['log', 'progress_bar', 'preds']}							)
											# Log results
											__lowercase								=   Path(pl_module.hparams.output_dir							)
											if type_path == "test":
															__lowercase								=   od / 'test_results.txt'
															__lowercase								=   od / 'test_generations.txt'
											else:
															# this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
															# If people want this it will be easy enough to add back.
															__lowercase								=   od / f"""{type_path}_results/{trainer.global_step:05d}.txt"""
															__lowercase								=   od / f"""{type_path}_generations/{trainer.global_step:05d}.txt"""
															results_file.parent.mkdir(exist_ok=_UpperCAmelCase							)
															generations_file.parent.mkdir(exist_ok=_UpperCAmelCase							)
											with open(_UpperCAmelCase	,					'a+'							) as writer:
															for key in sorted(_UpperCAmelCase							):
																			if key in ["log", "progress_bar", "preds"]:
																							continue
																			__lowercase								=   metrics[key]
																			if isinstance(_UpperCAmelCase	,					torch.Tensor							):
																							__lowercase								=   val.item()
																			__lowercase								=   f"""{key}: {val:.6f}\n"""
																			writer.write(_UpperCAmelCase							)

											if not save_generations:
															return

											if "preds" in metrics:
															__lowercase								=   '\n'.join(metrics['preds']							)
															generations_file.open('w+'							).write(_UpperCAmelCase							)





							@rank_zero_only
							def      a__					(					self    :  str	,					_UpperCAmelCase    :  Optional[int]	,					_UpperCAmelCase    :  str							)		->						Optional[int]:



											"""simple docstring"""


											try:
															__lowercase								=   pl_module.model.model.num_parameters()
											except AttributeError:
															__lowercase								=   pl_module.model.num_parameters()

											__lowercase								=   count_trainable_parameters(_UpperCAmelCase							)
											# mp stands for million parameters
											trainer.logger.log_metrics({'n_params': npars, 'mp': npars / 1e6, 'grad_mp': n_trainable_pars / 1e6}							)





							@rank_zero_only
							def      a__					(					self    :  Union[str, Any]	,					_UpperCAmelCase    :  pl.Trainer	,					_UpperCAmelCase    :  pl.LightningModule							)		->						Optional[Any]:



											"""simple docstring"""


											save_json(pl_module.metrics	,					pl_module.metrics_save_path							)
											return self._write_logs(_UpperCAmelCase	,					_UpperCAmelCase	,					'test'							)





							@rank_zero_only
							def      a__					(					self    :  int	,					_UpperCAmelCase    :  pl.Trainer	,					_UpperCAmelCase    :  Tuple							)		->						str:



											"""simple docstring"""


											save_json(pl_module.metrics	,					pl_module.metrics_save_path							)
											# Uncommenting this will save val generations
											# return self._write_logs(trainer, pl_module, "valid")






def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: List[Any]       ,				SCREAMING_SNAKE_CASE			: List[str]       ,				SCREAMING_SNAKE_CASE			: Union[str, Any]=1       ,				SCREAMING_SNAKE_CASE			: str=False     )	->		List[str]:
				if metric == "rouge2":
								__lowercase								=   '{val_avg_rouge2:.4f}-{step_count}'
				elif metric == "bleu":
								__lowercase								=   '{val_avg_bleu:.4f}-{step_count}'
				elif metric == "loss":
								__lowercase								=   '{val_avg_loss:.4f}-{step_count}'
				else:
								raise NotImplementedError(
								    F"""seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to"""
								    ' this function.'     )

				__lowercase								=   ModelCheckpoint(
				    dirpath=SCREAMING_SNAKE_CASE       ,				filename=SCREAMING_SNAKE_CASE       ,				monitor=F"""val_{metric}"""       ,				mode='min' if 'loss' in metric else 'max'       ,				save_top_k=SCREAMING_SNAKE_CASE       ,				)
				return checkpoint_callback




def 				__SCREAMING_SNAKE_CASE  (						SCREAMING_SNAKE_CASE			: Any       ,				SCREAMING_SNAKE_CASE			: int     )	->		Union[str, Any]:
				return EarlyStopping(
				    monitor=F"""val_{metric}"""       ,				mode='min' if 'loss' in metric else 'max'       ,				patience=SCREAMING_SNAKE_CASE       ,				verbose=SCREAMING_SNAKE_CASE       ,				)





