






import argparse
import gc
import os
import threading

import evaluate
import psutil
import torch
from datasets import load_dataset
from torch.distributed.fsdp.fully_sharded_data_parallel import FullOptimStateDictConfig, FullStateDictConfig
from torch.utils.data import DataLoader
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    get_linear_schedule_with_warmup,
    set_seed,
)

from accelerate import Accelerator, DistributedType, FullyShardedDataParallelPlugin


########################################################################
# This is a fully working simple example to use Accelerate
#
# This example trains a Bert base model on GLUE MRPC
# in any of the following settings (with the same script):
#   - single CPU or single GPU
#   - multi GPUS (using PyTorch distributed mode)
#   - (multi) TPUs
#   - fp16 (mixed-precision) or fp32 (normal precision)
#   - FSDP
#
# This example also demonstrates the checkpointing and sharding capabilities
#
# To run it in each of these various modes, follow the instructions
# in the readme for examples:
# https://github.com/huggingface/accelerate/tree/main/examples
#
########################################################################


_lowercase	:   List[Any]    	=16
_lowercase	:   int    	=32

def 			lowerCAmelCase_  (						_lowercase       :     List[Any])							-> List[str]:







			"""simple docstring"""
			return int(x / 2**20)

class      snake_case__       :







						"""simple docstring"""




						def __enter__(     self						)      ->   Optional[Any]:
									"""simple docstring"""




									gc.collect()
									torch.cuda.empty_cache()
									torch.cuda.reset_max_memory_allocated()  # reset the peak gauge to zero
									a__     :		Optional[Any]		 =						torch.cuda.memory_allocated()
									a__     :		int		 =						psutil.Process()

									a__     :		List[str]		 =						self.cpu_mem_used()
									a__     :		Dict		 =						True
									a__     :		Optional[Any]		 =						threading.Thread(target=self.peak_monitor_func						)
									a__     :		Dict		 =						True
									peak_monitor_thread.start()
									return self

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Dict:
									"""simple docstring"""




									return self.process.memory_info().rss

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   int:
									"""simple docstring"""




									a__     :		Union[str, Any]		 =						-1

									while True:
												a__     :		str		 =						max(self.cpu_mem_used()      ,     self.cpu_peak						)

												# can't sleep or will not catch the peak right (this comment is here on purpose)
												# time.sleep(0.001) # 1msec

												if not self.peak_monitoring:
															break

						def __exit__(     self      ,     *__lowercase						)      ->   Optional[Any]:
									"""simple docstring"""




									a__     :		Tuple		 =						False

									gc.collect()
									torch.cuda.empty_cache()
									a__     :		Union[str, Any]		 =						torch.cuda.memory_allocated()
									a__     :		List[Any]		 =						torch.cuda.max_memory_allocated()
									a__     :		Tuple		 =						bamb(self.end - self.begin						)
									a__     :		Optional[int]		 =						bamb(self.peak - self.begin						)

									a__     :		List[Any]		 =						self.cpu_mem_used()
									a__     :		Dict		 =						bamb(self.cpu_end - self.cpu_begin						)
									a__     :		Optional[Any]		 =						bamb(self.cpu_peak - self.cpu_begin						)
									# print(f"delta used/peak {self.used:4d}/{self.peaked:4d}")


# For testing only
if os.environ.get("TESTING_MOCKED_DATALOADERS", None) == "1":
		from accelerate.test_utils.training import mocked_dataloaders

		_lowercase	:   int    	=mocked_dataloaders  # noqa: F811

def 			lowerCAmelCase_  (						_lowercase       :     Union[str, Any]		,     _lowercase       :     Dict)							-> List[Any]:







			"""simple docstring"""






			if os.environ.get("""TESTING_MOCKED_DATALOADERS"""		,     _UpperCAmelCase) == "1":
						a__     :		Optional[int]		 =						2

			# New Code #
			# Pass the advanced FSDP settings not part of the accelerate config by creating fsdp_plugin
			a__     :		List[str]		 =						FullyShardedDataParallelPlugin(
			    state_dict_config=FullStateDictConfig(offload_to_cpu=_UpperCAmelCase		,     ranka_only=_UpperCAmelCase)		,     optim_state_dict_config=FullOptimStateDictConfig(offload_to_cpu=_UpperCAmelCase		,     ranka_only=_UpperCAmelCase)		,     )

			# Initialize accelerator
			if args.with_tracking:
						a__     :		Any		 =						Accelerator(
						    cpu=args.cpu		,     mixed_precision=args.mixed_precision		,     log_with="""wandb"""		,     project_dir=args.logging_dir		,     fsdp_plugin=_UpperCAmelCase		,     )
			else:
						a__     :		Optional[int]		 =						Accelerator(fsdp_plugin=_UpperCAmelCase)
			accelerator.print(accelerator.distributed_type)

			if hasattr(args.checkpointing_steps		,     """isdigit"""):
						if args.checkpointing_steps == "epoch":
									a__     :		List[Any]		 =						args.checkpointing_steps
						elif args.checkpointing_steps.isdigit():
									a__     :		Optional[int]		 =						int(args.checkpointing_steps)
						else:
									raise ValueError(
									    F'''Argument `checkpointing_steps` must be either a number or `epoch`. `{args.checkpointing_steps}` passed.''')
			else:
						a__     :		List[Any]		 =						None
			# Sample hyper-parameters for learning rate, batch size, seed and a few other HPs
			a__     :		Union[str, Any]		 =						config['lr']
			a__     :		Union[str, Any]		 =						int(config["""num_epochs"""])
			a__     :		Tuple		 =						int(config["""seed"""])
			a__     :		Any		 =						int(config["""batch_size"""])

			# We need to initialize the trackers we use, and also store our configuration
			if args.with_tracking:
						a__     :		Dict		 =						vars(_UpperCAmelCase)
						accelerator.init_trackers("""fsdp_glue_no_trainer"""		,     _UpperCAmelCase)

			a__     :		List[str]		 =						AutoTokenizer.from_pretrained(args.model_name_or_path)
			a__     :		Any		 =						load_dataset("""glue"""		,     """mrpc""")
			a__     :		Any		 =						evaluate.load("""glue"""		,     """mrpc""")

			def tokenize_function(_lowercase       :     str):
						# max_length=None => use the model max length (it's actually the default)
						a__     :		int		 =						tokenizer(examples["""sentence1"""]		,     examples["""sentence2"""]		,     truncation=_UpperCAmelCase		,     max_length=_UpperCAmelCase)
						return outputs

			# Apply the method we just defined to all the examples in all the splits of the dataset
			# starting with the main process first:
			with accelerator.main_process_first():
						a__     :		Union[str, Any]		 =						datasets.map(
						    _UpperCAmelCase		,     batched=_UpperCAmelCase		,     remove_columns=["""idx""", """sentence1""", """sentence2"""]		,     )

			# We also rename the 'label' column to 'labels' which is the expected name for labels by the models of the
			# transformers library
			a__     :		Union[str, Any]		 =						tokenized_datasets.rename_column("""label"""		,     """labels""")

			# If the batch size is too big we use gradient accumulation
			a__     :		Optional[Any]		 =						1
			if batch_size > MAX_GPU_BATCH_SIZE and accelerator.distributed_type != DistributedType.TPU:
						a__     :		str		 =						batch_size // MAX_GPU_BATCH_SIZE
						a__     :		Optional[int]		 =						MAX_GPU_BATCH_SIZE

			def collate_fn(_lowercase       :     Tuple):
						# On TPU it's best to pad everything to the same length or training will be very slow.
						a__     :		List[str]		 =						128 if accelerator.distributed_type == DistributedType.TPU else None
						# When using mixed precision we want round multiples of 8/16
						if accelerator.mixed_precision == "fp8":
									a__     :		Tuple		 =						16
						elif accelerator.mixed_precision != "no":
									a__     :		List[Any]		 =						8
						else:
									a__     :		Union[str, Any]		 =						None

						return tokenizer.pad(
						    _UpperCAmelCase		,     padding="""longest"""		,     max_length=_UpperCAmelCase		,     pad_to_multiple_of=_UpperCAmelCase		,     return_tensors="""pt"""		,     )

			# Instantiate dataloaders.
			a__     :		int		 =						DataLoader(
			    tokenized_datasets["""train"""]		,     shuffle=_UpperCAmelCase		,     collate_fn=_UpperCAmelCase		,     batch_size=_UpperCAmelCase)
			a__     :		int		 =						DataLoader(
			    tokenized_datasets["""validation"""]		,     shuffle=_UpperCAmelCase		,     collate_fn=_UpperCAmelCase		,     batch_size=_UpperCAmelCase)

			set_seed(_UpperCAmelCase)

			# Instantiate the model (we build the model here so that the seed also control new weights initialization)
			a__     :		Optional[int]		 =						AutoModelForSequenceClassification.from_pretrained(
			    args.model_name_or_path		,     return_dict=_UpperCAmelCase		,     low_cpu_mem_usage=_UpperCAmelCase)

			# New Code #
			# For FSDP feature, it is highly recommended and efficient to prepare the model before creating optimizer
			a__     :		int		 =						accelerator.prepare(_UpperCAmelCase)
			accelerator.print(_UpperCAmelCase)

			# Instantiate optimizer
			# New Code #
			# For FSDP feature, at present it doesn't support multiple parameter groups,
			# so we need to create a single parameter group for the whole model
			a__     :		str		 =						torch.optim.AdamW(params=model.parameters()		,     lr=_UpperCAmelCase		,     weight_decay=2e-4)

			# Instantiate scheduler
			a__     :		Optional[Any]		 =						get_linear_schedule_with_warmup(
			    optimizer=_UpperCAmelCase		,     num_warmup_steps=10		,     num_training_steps=(len(_UpperCAmelCase) * num_epochs) // gradient_accumulation_steps		,     )

			# New Code #
			# For FSDP feature, prepare everything except the model as we have already prepared the model
			# before creating the optimizer
			# There is no specific order to remember, we just need to unpack the objects in the same order we gave them to the
			# prepare method.
			a__     :		Any		 =						accelerator.prepare(
			    _UpperCAmelCase		,     _UpperCAmelCase		,     _UpperCAmelCase		,     _UpperCAmelCase)

			a__     :		Optional[Any]		 =						0

			# Potentially load in the weights and states from a previous save
			if args.resume_from_checkpoint:
						if args.resume_from_checkpoint is not None or args.resume_from_checkpoint != "":
									accelerator.print(F'''Resumed from checkpoint: {args.resume_from_checkpoint}''')
									accelerator.load_state(args.resume_from_checkpoint)
									a__     :		Optional[Any]		 =						os.path.basename(args.resume_from_checkpoint)
						else:
									# Get the most recent checkpoint
									a__     :		Tuple		 =						[f.name for f in os.scandir(os.getcwd()) if f.is_dir()]
									dirs.sort(key=os.path.getctime)
									a__     :		Optional[int]		 =						dirs[-1]  # Sorts folders by date modified, most recent checkpoint is the last
						# Extract `epoch_{i}` or `step_{i}`
						a__     :		str		 =						os.path.splitext(_UpperCAmelCase)[0]

						if "epoch" in training_difference:
									num_epochs -= int(training_difference.replace("""epoch_"""		,     """"""))
									a__     :		str		 =						None
						else:
									a__     :		Optional[int]		 =						int(training_difference.replace("""step_"""		,     """"""))
									num_epochs -= resume_step // len(_UpperCAmelCase)
									# If resuming by step, we also need to know exactly how far into the DataLoader we went
									a__     :		str		 =						(num_epochs * len(_UpperCAmelCase)) - resume_step

    # Now we train the model
			for epoch in range(_UpperCAmelCase):
						# New Code #
						# context manager to track the peak memory usage during the training epoch
						with TorchTracemalloc() as tracemalloc:
									model.train()
									if args.with_tracking:
												a__     :		Optional[Any]		 =						0
									for step, batch in enumerate(_UpperCAmelCase):
												# We need to skip steps until we reach the resumed step
												if args.resume_from_checkpoint and epoch == 0:
															if resume_step is not None and step < resume_step:
																		pass
                # We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device)
												a__     :		List[Any]		 =						model(**_UpperCAmelCase)
												a__     :		str		 =						outputs.loss
												a__     :		List[Any]		 =						loss / gradient_accumulation_steps
												# We keep track of the loss at each epoch
												if args.with_tracking:
															total_loss += loss.detach().float()
												accelerator.backward(_UpperCAmelCase)
												if step % gradient_accumulation_steps == 0:
															optimizer.step()
															lr_scheduler.step()
															optimizer.zero_grad()
															# accelerator.print(lr_scheduler.get_lr())

												overall_step += 1

												if isinstance(_UpperCAmelCase		,     _UpperCAmelCase):
															a__     :		Optional[int]		 =						F'''step_{overall_step}'''
															if overall_step % checkpointing_steps == 0:
																		if args.output_dir is not None:
																					a__     :		List[str]		 =						os.path.join(args.output_dir		,     _UpperCAmelCase)
																		accelerator.save_state(_UpperCAmelCase)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
						accelerator.print("""Memory before entering the train : {}""".format(bamb(tracemalloc.begin)))
						accelerator.print("""Memory consumed at the end of the train (end-begin): {}""".format(tracemalloc.used))
						accelerator.print("""Peak Memory consumed during the train (max-begin): {}""".format(tracemalloc.peaked))
						accelerator.print(
						    """Total Peak Memory consumed during the train (max): {}""".format(
						        tracemalloc.peaked + bamb(tracemalloc.begin)))
						# Logging the peak memory usage of the GPU to the tracker
						if args.with_tracking:
									accelerator.log(
									    {
									        """train_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin),
									    }		,     step=_UpperCAmelCase		,     )

						# New Code #
						# context manager to track the peak memory usage during the evaluation
						with TorchTracemalloc() as tracemalloc:
									model.eval()
									for step, batch in enumerate(_UpperCAmelCase):
												# We could avoid this line since we set the accelerator with `device_placement=True`.
												batch.to(accelerator.device)
												with torch.no_grad():
															a__     :		List[Any]		 =						model(**_UpperCAmelCase)
												a__     :		Optional[Any]		 =						outputs.logits.argmax(dim=-1)
												a__     :		Union[str, Any]		 =						accelerator.gather_for_metrics((predictions, batch["""labels"""]))
												metric.add_batch(
												    predictions=_UpperCAmelCase		,     references=_UpperCAmelCase		,     )

									a__     :		Dict		 =						metric.compute()
									# Use accelerator.print to print only on the main process.
									accelerator.print(F'''epoch {epoch}:'''		,     _UpperCAmelCase)
									if args.with_tracking:
												accelerator.log(
												    {
												        """accuracy""": eval_metric["""accuracy"""],
												        """f1""": eval_metric["""f1"""],
												        """train_loss""": total_loss.item() / len(_UpperCAmelCase),
												    }		,     step=_UpperCAmelCase		,     )

									if checkpointing_steps == "epoch":
												a__     :		str		 =						F'''epoch_{epoch}'''
												if args.output_dir is not None:
															a__     :		List[str]		 =						os.path.join(args.output_dir		,     _UpperCAmelCase)
												accelerator.save_state(_UpperCAmelCase)
        # New Code #
        # Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage
						accelerator.print("""Memory before entering the eval : {}""".format(bamb(tracemalloc.begin)))
						accelerator.print("""Memory consumed at the end of the eval (end-begin): {}""".format(tracemalloc.used))
						accelerator.print("""Peak Memory consumed during the eval (max-begin): {}""".format(tracemalloc.peaked))
						accelerator.print(
						    """Total Peak Memory consumed during the eval (max): {}""".format(tracemalloc.peaked + bamb(tracemalloc.begin)))
						# Logging the peak memory usage of the GPU to the tracker
						if args.with_tracking:
									accelerator.log(
									    {
									        """eval_total_peak_memory""": tracemalloc.peaked + bamb(tracemalloc.begin),
									    }		,     step=_UpperCAmelCase		,     )

			if args.with_tracking:
						accelerator.end_training()

def 			lowerCAmelCase_  (						)							-> Union[str, Any]:







			"""simple docstring"""
			a__     :		Dict		 =						argparse.ArgumentParser(description="""Simple example of training script.""")
			parser.add_argument(
			    """--mixed_precision"""		,     type=_UpperCAmelCase		,     default=_UpperCAmelCase		,     choices=["""no""", """fp16""", """bf16""", """fp8"""]		,     help="""Whether to use mixed precision. Choose"""
			    """between fp16 and bf16 (bfloat16). Bf16 requires PyTorch >= 1.10."""
			    """and an Nvidia Ampere GPU."""		,     )
			parser.add_argument("""--cpu"""		,     action="""store_true"""		,     help="""If passed, will train on the CPU.""")
			parser.add_argument(
			    """--checkpointing_steps"""		,     type=_UpperCAmelCase		,     default=_UpperCAmelCase		,     help="""Whether the various states should be saved at the end of every n steps, or \'epoch\' for each epoch."""		,     )
			parser.add_argument(
			    """--resume_from_checkpoint"""		,     type=_UpperCAmelCase		,     default=_UpperCAmelCase		,     help="""If the training should continue from a checkpoint folder."""		,     )
			parser.add_argument(
			    """--with_tracking"""		,     action="""store_true"""		,     help="""Whether to load in all available experiment trackers from the environment and use them for logging."""		,     )
			parser.add_argument(
			    """--output_dir"""		,     type=_UpperCAmelCase		,     default="""."""		,     help="""Optional save directory where all checkpoint folders will be stored. Default is the current working directory."""		,     )
			parser.add_argument(
			    """--logging_dir"""		,     type=_UpperCAmelCase		,     default="""logs"""		,     help="""Location on where to store experiment tracking logs`"""		,     )
			parser.add_argument(
			    """--model_name_or_path"""		,     type=_UpperCAmelCase		,     help="""Path to pretrained model or model identifier from huggingface.co/models."""		,     required=_UpperCAmelCase		,     )
			a__     :		Union[str, Any]		 =						parser.parse_args()
			a__     :		Any		 =						{'lr': 2e-5, 'num_epochs': 3, 'seed': 42, 'batch_size': 16}
			training_function(_UpperCAmelCase		,     _UpperCAmelCase)


if __name__ == "__main__":
		main()









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







						"""simple docstring"""




						def __init__(     self      ,     __lowercase      ,     )      ->   str:
									"""simple docstring"""




									a__     :		Union[str, Any]		 =						parent
									a__     :		Dict		 =						1_3
									a__     :		List[Any]		 =						7
									a__     :		Tuple		 =						True
									a__     :		List[str]		 =						True
									a__     :		List[str]		 =						True
									a__     :		Optional[int]		 =						True
									a__     :		int		 =						True
									a__     :		Optional[int]		 =						9_9
									a__     :		Optional[int]		 =						3_2
									a__     :		Tuple		 =						2
									a__     :		Any		 =						4
									a__     :		Optional[Any]		 =						3_7
									a__     :		Union[str, Any]		 =						"""gelu"""
									a__     :		Dict		 =						0.1
									a__     :		Optional[int]		 =						0.1
									a__     :		int		 =						5_1_2
									a__     :		Any		 =						1_6
									a__     :		List[Any]		 =						2
									a__     :		Optional[Any]		 =						0.0_2
									a__     :		List[Any]		 =						3
									a__     :		Dict		 =						4
									a__     :		List[str]		 =						None
									a__     :		Tuple		 =						self.vocab_size - 1

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Dict:
									"""simple docstring"""




									a__     :		Tuple		 =						ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size						)

									a__     :		int		 =						None
									if self.use_input_mask:
												a__     :		Optional[int]		 =						random_attention_mask([self.batch_size, self.seq_length]						)

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

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

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

									a__     :		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      ,     )

									a__     :		Optional[int]		 =						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 							SCREAMING_SNAKE_CASE__(     self      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     *__lowercase						)      ->   str:
									"""simple docstring"""




									a__     :		Dict		 =						TFOpenAIGPTModel(config=__lowercase						)
									a__     :		Tuple		 =						{"""input_ids""": input_ids, """attention_mask""": input_mask, """token_type_ids""": token_type_ids}
									a__     :		int		 =						model(__lowercase						)

									a__     :		List[Any]		 =						[input_ids, input_mask]
									a__     :		int		 =						model(__lowercase						)

									a__     :		Optional[int]		 =						model(__lowercase						)

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

						def 							SCREAMING_SNAKE_CASE__(     self      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     *__lowercase						)      ->   int:
									"""simple docstring"""




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

						def 							SCREAMING_SNAKE_CASE__(     self      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     *__lowercase						)      ->   Tuple:
									"""simple docstring"""




									a__     :		List[Any]		 =						TFOpenAIGPTDoubleHeadsModel(config=__lowercase						)

									a__     :		int		 =						tf.tile(tf.expand_dims(__lowercase      ,     1						)      ,     (1, self.num_choices, 1)						)
									a__     :		List[str]		 =						tf.tile(tf.expand_dims(__lowercase      ,     1						)      ,     (1, self.num_choices, 1)						)
									a__     :		Dict		 =						tf.tile(tf.expand_dims(__lowercase      ,     1						)      ,     (1, self.num_choices, 1)						)

									a__     :		Optional[Any]		 =						{
									    """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,
									}
									a__     :		Optional[Any]		 =						model(__lowercase						)
									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 							SCREAMING_SNAKE_CASE__(     self      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     *__lowercase						)      ->   Dict:
									"""simple docstring"""




									a__     :		Tuple		 =						self.num_labels
									a__     :		Optional[int]		 =						ids_tensor([self.batch_size]      ,     self.type_sequence_label_size						)
									a__     :		Tuple		 =						{
									    """input_ids""": input_ids,
									    """attention_mask""": input_mask,
									    """token_type_ids""": token_type_ids,
									    """labels""": sequence_labels,
									}
									a__     :		Optional[Any]		 =						TFOpenAIGPTForSequenceClassification(__lowercase						)
									a__     :		Union[str, Any]		 =						model(__lowercase						)
									self.parent.assertEqual(result.logits.shape      ,     (self.batch_size, self.num_labels)						)

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   int:
									"""simple docstring"""




									a__     :		Optional[Any]		 =						self.prepare_config_and_inputs()

									(
									    (
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		(
									    a__
									) ,		
									)     :		Dict		 =						config_and_inputs

									a__     :		Union[str, Any]		 =						{"""input_ids""": input_ids, """token_type_ids""": token_type_ids, """attention_mask""": input_mask}
									return config, inputs_dict




@require_tf
class      snake_case__       (A__      ,			A__      ,			unittest.TestCase	):







						"""simple docstring"""




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

						def 							SCREAMING_SNAKE_CASE__(     self      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase      ,     __lowercase						)      ->   Optional[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 							SCREAMING_SNAKE_CASE__(     self						)      ->   int:
									"""simple docstring"""




									a__     :		Optional[Any]		 =						TFOpenAIGPTModelTester(self						)
									a__     :		List[Any]		 =						ConfigTester(self      ,     config_class=__lowercase      ,     n_embd=3_7						)

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Optional[Any]:
									"""simple docstring"""




									self.config_tester.run_common_tests()

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Union[str, Any]:
									"""simple docstring"""




									a__     :		Tuple		 =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_model(*__lowercase						)

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Optional[int]:
									"""simple docstring"""




									a__     :		Optional[Any]		 =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_lm_head(*__lowercase						)

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   List[str]:
									"""simple docstring"""




									a__     :		int		 =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_double_head(*__lowercase						)

						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Dict:
									"""simple docstring"""




									a__     :		Optional[Any]		 =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_openai_gpt_for_sequence_classification(*__lowercase						)

						@slow
						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Optional[Any]:
									"""simple docstring"""




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




@require_tf
class      snake_case__       (unittest.TestCase	):







						"""simple docstring"""




						@slow
						def 							SCREAMING_SNAKE_CASE__(     self						)      ->   Dict:
									"""simple docstring"""




									a__     :		int		 =						TFOpenAIGPTLMHeadModel.from_pretrained("""openai-gpt"""						)
									a__     :		Tuple		 =						tf.convert_to_tensor([[4_8_1, 4_7_3_5, 5_4_4]]      ,     dtype=tf.intaa						)  # the president is
									a__     :		List[str]		 =						[
									    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

									a__     :		Union[str, Any]		 =						model.generate(__lowercase      ,     do_sample=__lowercase						)
									self.assertListEqual(output_ids[0].numpy().tolist()      ,     __lowercase						)




