

'''simple docstring'''


import tempfile
import unittest

from transformers import TaConfig, is_torch_available
from transformers.testing_utils import (
    require_sentencepiece,
    require_tokenizers,
    require_torch,
    slow,
    torch_device,
)

from ...generation.test_utils import GenerationTesterMixin
from ...test_modeling_common import ModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import AutoTokenizer, UMTaForConditionalGeneration, UMTaForQuestionAnswering, UMTaModel



class 					lowercase						:


							def __init__(			self    ,			_snake_case    ,			_snake_case=99    ,			_snake_case=13    ,			_snake_case=7    ,			_snake_case=9    ,			_snake_case=True    ,			_snake_case=True    ,			_snake_case=False    ,			_snake_case=32    ,			_snake_case=5    ,			_snake_case=4    ,			_snake_case=37    ,			_snake_case=8    ,			_snake_case=0.1    ,			_snake_case=0.002    ,			_snake_case=1    ,			_snake_case=0    ,			_snake_case=0    ,			_snake_case=None    ,			_snake_case=None    ,			)							->				Union[str, Any]:
								UpperCAmelCase_		:      Dict       	=					parent
								UpperCAmelCase_		:      List[str]       	=					batch_size
								UpperCAmelCase_		:      Union[str, Any]       	=					encoder_seq_length
								UpperCAmelCase_		:      Tuple       	=					decoder_seq_length
								# For common tests
								UpperCAmelCase_		:      Optional[Any]       	=					self.decoder_seq_length
								UpperCAmelCase_		:      int       	=					is_training
								UpperCAmelCase_		:      str       	=					use_attention_mask
								UpperCAmelCase_		:      int       	=					use_labels
								UpperCAmelCase_		:      Optional[Any]       	=					vocab_size
								UpperCAmelCase_		:      Tuple       	=					hidden_size
								UpperCAmelCase_		:      Optional[int]       	=					num_hidden_layers
								UpperCAmelCase_		:      Union[str, Any]       	=					num_attention_heads
								UpperCAmelCase_		:      Optional[int]       	=					d_ff
								UpperCAmelCase_		:      Dict       	=					relative_attention_num_buckets
								UpperCAmelCase_		:      Optional[Any]       	=					dropout_rate
								UpperCAmelCase_		:      Any       	=					initializer_factor
								UpperCAmelCase_		:      Optional[int]       	=					eos_token_id
								UpperCAmelCase_		:      Any       	=					pad_token_id
								UpperCAmelCase_		:      Optional[Any]       	=					decoder_start_token_id
								UpperCAmelCase_		:      Union[str, Any]       	=					None
								UpperCAmelCase_		:      Dict       	=					decoder_layers


							def 						_snake_case							(			self)							->				Optional[int]:
								return TaConfig.from_pretrained('google/umt5-base')


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case=None    ,			_snake_case=None    ,			_snake_case=None    ,			_snake_case=None    ,			_snake_case=None    ,			)							->				Any:
								if attention_mask is None:
									UpperCAmelCase_		:      List[str]       	=					input_ids.ne(config.pad_token_id)
								if decoder_attention_mask is None:
									UpperCAmelCase_		:      Optional[Any]       	=					decoder_input_ids.ne(config.pad_token_id)
								if head_mask is None:
									UpperCAmelCase_		:      Union[str, Any]       	=					torch.ones(config.num_hidden_layers    ,			config.num_attention_heads    ,			device=_snake_case)
								if decoder_head_mask is None:
									UpperCAmelCase_		:      Tuple       	=					torch.ones(config.num_decoder_layers    ,			config.num_attention_heads    ,			device=_snake_case)
								if cross_attn_head_mask is None:
									UpperCAmelCase_		:      int       	=					torch.ones(
									    config.num_decoder_layers    ,			config.num_attention_heads    ,			device=_snake_case)
								return {
								    "input_ids": input_ids,
								    "decoder_input_ids": decoder_input_ids,
								    "attention_mask": attention_mask,
								    "decoder_attention_mask": decoder_attention_mask,
								    "head_mask": head_mask,
								    "decoder_head_mask": decoder_head_mask,
								    "cross_attn_head_mask": cross_attn_head_mask,
								}


							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_		:      Optional[Any]       	=					ids_tensor([self.batch_size, self.encoder_seq_length]    ,			self.vocab_size)
								UpperCAmelCase_		:      Any       	=					ids_tensor([self.batch_size, self.decoder_seq_length]    ,			self.vocab_size)

								# we need to clamp the input ids here to avoid having pad token in between
								# this is because for NllbMoe the position_ids are prepared such that
								# all pad tokens have pos id = 2 and rest are between 2..seq_length
								# and the seq_length here is seq_length - num_pad_tokens
								# but when using past, there is no way of knowing if the past input ids had
								# pad tokens in them, which results in incorrect seq_lenth and which in turn results in
								# position_ids being off by num_pad_tokens in past input
								UpperCAmelCase_		:      Optional[int]       	=					input_ids.clamp(self.pad_token_id + 1)
								UpperCAmelCase_		:      Dict       	=					decoder_input_ids.clamp(self.pad_token_id + 1)

								UpperCAmelCase_		:      List[Any]       	=					self.get_config()
								UpperCAmelCase_		:      Any       	=					config.num_attention_heads
								UpperCAmelCase_		:      List[str]       	=					self.prepare_inputs_dict(_snake_case    ,			_snake_case    ,			_snake_case)
								return config, input_dict


							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_	,    UpperCAmelCase_		:      List[Any]       	=					self.prepare_config_and_inputs()
								return config, inputs_dict


							def 						_snake_case							(			self)							->				str:
								return TaConfig(
								    vocab_size=166    ,			d_model=self.hidden_size    ,			d_ff=self.d_ff    ,			d_kv=self.hidden_size // self.num_attention_heads    ,			num_layers=self.num_hidden_layers    ,			num_decoder_layers=self.decoder_layers    ,			num_heads=self.num_attention_heads    ,			relative_attention_num_buckets=self.relative_attention_num_buckets    ,			dropout_rate=self.dropout_rate    ,			initializer_factor=self.initializer_factor    ,			eos_token_id=self.eos_token_id    ,			bos_token_id=self.pad_token_id    ,			pad_token_id=self.pad_token_id    ,			decoder_start_token_id=self.decoder_start_token_id    ,			)


							def 						_snake_case							(			self)							->				List[str]:
								return TaConfig(
								    vocab_size=self.vocab_size    ,			d_model=self.hidden_size    ,			d_ff=self.d_ff    ,			d_kv=self.hidden_size // self.num_attention_heads    ,			num_layers=self.num_hidden_layers    ,			num_decoder_layers=self.decoder_layers    ,			num_heads=self.num_attention_heads    ,			relative_attention_num_buckets=self.relative_attention_num_buckets    ,			dropout_rate=self.dropout_rate    ,			initializer_factor=self.initializer_factor    ,			eos_token_id=self.eos_token_id    ,			bos_token_id=self.pad_token_id    ,			pad_token_id=self.pad_token_id    ,			decoder_start_token_id=self.decoder_start_token_id    ,			)


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				str:
								UpperCAmelCase_		:      str       	=					UMTaModel(config=_snake_case)
								model.to(_snake_case)
								model.eval()
								UpperCAmelCase_		:      Optional[Any]       	=					model(
								    input_ids=_snake_case    ,			decoder_input_ids=_snake_case    ,			attention_mask=_snake_case    ,			decoder_attention_mask=_snake_case    ,			)
								UpperCAmelCase_		:      Tuple       	=					model(input_ids=_snake_case    ,			decoder_input_ids=_snake_case)
								UpperCAmelCase_		:      List[str]       	=					result.last_hidden_state
								UpperCAmelCase_		:      Union[str, Any]       	=					result.past_key_values
								UpperCAmelCase_		:      Optional[Any]       	=					result.encoder_last_hidden_state

								self.parent.assertEqual(encoder_output.size()    ,			(self.batch_size, self.encoder_seq_length, self.hidden_size))
								self.parent.assertEqual(decoder_output.size()    ,			(self.batch_size, self.decoder_seq_length, self.hidden_size))
								# There should be `num_layers` key value embeddings stored in decoder_past
								self.parent.assertEqual(len(_snake_case)    ,			config.num_layers)
								# There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past tuple
								self.parent.assertEqual(len(decoder_past[0])    ,			4)


							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case    ,			)							->				Optional[Any]:
								UpperCAmelCase_		:      Dict       	=					UMTaModel(config=_snake_case).get_decoder().to(_snake_case).eval()
								# first forward pass
								UpperCAmelCase_		:      Optional[Any]       	=					model(_snake_case    ,			use_cache=_snake_case)
								UpperCAmelCase_		:      List[Any]       	=					model(_snake_case)
								UpperCAmelCase_		:      List[str]       	=					model(_snake_case    ,			use_cache=_snake_case)

								self.parent.assertTrue(len(_snake_case) == len(_snake_case))
								self.parent.assertTrue(len(_snake_case) == len(_snake_case) + 1)

								UpperCAmelCase_	,    UpperCAmelCase_		:      Optional[Any]       	=					outputs.to_tuple()

								# create hypothetical next token and extent to next_input_ids
								UpperCAmelCase_		:      Union[str, Any]       	=					ids_tensor((self.batch_size, 1)    ,			config.vocab_size)

								# append to next input_ids and
								UpperCAmelCase_		:      Any       	=					torch.cat([input_ids, next_tokens]    ,			dim=-1)

								UpperCAmelCase_		:      Tuple       	=					model(_snake_case)['last_hidden_state']
								UpperCAmelCase_		:      Tuple       	=					model(_snake_case    ,			past_key_values=_snake_case)['last_hidden_state']

								# select random slice
								UpperCAmelCase_		:      Dict       	=					ids_tensor((1,)    ,			output_from_past.shape[-1]).item()
								UpperCAmelCase_		:      Any       	=					output_from_no_past[:, -1, random_slice_idx].detach()
								UpperCAmelCase_		:      int       	=					output_from_past[:, 0, random_slice_idx].detach()

								# test that outputs are equal for slice
								self.parent.assertTrue(torch.allclose(_snake_case    ,			_snake_case    ,			atol=1e-3))





							def 						_snake_case							(			self    ,			_snake_case    ,			_snake_case    ,			)							->				Optional[Any]:
								UpperCAmelCase_		:      Tuple       	=					UMTaModel(config=_snake_case).to(_snake_case).half().eval()
								UpperCAmelCase_		:      Optional[int]       	=					model(**_snake_case)['last_hidden_state']
								self.parent.assertFalse(torch.isnan(_snake_case).any().item())



@require_torch
class 					lowercase						(  a_,					a_,					a_,					unittest.TestCase ):
							_lowerCamelCase		:      Union[str, Any]=   (
							    (UMTaModel, UMTaForConditionalGeneration, UMTaForQuestionAnswering) if is_torch_available() else ()
							)
							_lowerCamelCase		:      Optional[Any]=   (UMTaForConditionalGeneration,) if is_torch_available() else ()
							_lowerCamelCase		:      str=   (
							    {
							        "conversational": UMTaForConditionalGeneration,
							        "feature-extraction": UMTaModel,
							        "summarization": UMTaForConditionalGeneration,
							        "text2text-generation": UMTaForConditionalGeneration,
							        "translation": UMTaForConditionalGeneration,
							        "question-answering": UMTaForQuestionAnswering,
							    }
							    if is_torch_available()
							    else {}
							)
							_lowerCamelCase		:      int=   True
							_lowerCamelCase		:      Dict=   False
							_lowerCamelCase		:      List[str]=   False
							_lowerCamelCase		:      Any=   True
							_lowerCamelCase		:      Union[str, Any]=   True
							# The small UMT5 model needs higher percentages for CPU/MP tests
							_lowerCamelCase		:      str=   [0.8, 0.9]


							def 						_snake_case							(			self)							->				List[Any]:
								UpperCAmelCase_		:      List[Any]       	=					UMTaModelTester(self)


							@unittest.skip('Test has a segmentation fault on torch 1.8.0')
							def 						_snake_case							(			self)							->				str:
								UpperCAmelCase_		:      Optional[Any]       	=					self.model_tester.prepare_config_and_inputs()
								UpperCAmelCase_		:      List[Any]       	=					UMTaModel(config_and_inputs[0]).to(_snake_case)
								with tempfile.TemporaryDirectory() as tmpdirname:
									torch.onnx.export(
									    _snake_case    ,			(config_and_inputs[1], config_and_inputs[3], config_and_inputs[2])    ,			F"""{tmpdirname}/t5_test.onnx"""    ,			export_params=_snake_case    ,			opset_version=9    ,			input_names=['input_ids', 'decoder_input_ids']    ,			)


							@unittest.skipIf(torch_device == 'cpu'    ,			'Cant do half precision')
							def 						_snake_case							(			self)							->				Optional[int]:
								UpperCAmelCase_		:      Dict       	=					self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model_fpaa_forward(*_snake_case)


							def 						_snake_case							(			self)							->				Optional[Any]:
								UpperCAmelCase_		:      Optional[Any]       	=					['encoder_attentions', 'decoder_attentions', 'cross_attentions']
								UpperCAmelCase_		:      str       	=					self.model_tester.prepare_config_and_inputs()
								UpperCAmelCase_		:      int       	=					config_and_inputs[0]
								UpperCAmelCase_		:      Optional[Any]       	=					UMTaForConditionalGeneration(_snake_case).eval()
								model.to(_snake_case)

								UpperCAmelCase_		:      Dict       	=					{
								    'head_mask': torch.zeros(config.num_layers    ,			config.num_heads    ,			device=_snake_case),
								    'decoder_head_mask': torch.zeros(config.num_decoder_layers    ,			config.num_heads    ,			device=_snake_case),
								    'cross_attn_head_mask': torch.zeros(config.num_decoder_layers    ,			config.num_heads    ,			device=_snake_case),
								}

								for attn_name, (name, mask) in zip(_snake_case    ,			head_masking.items()):
									UpperCAmelCase_		:      str       	=					{name: mask}
									# Explicitly pass decoder_head_mask as it is required from T5 model when head_mask specified
									if name == "head_mask":
										UpperCAmelCase_		:      int       	=					torch.ones(
										    config.num_decoder_layers    ,			config.num_heads    ,			device=_snake_case)

									UpperCAmelCase_		:      Optional[int]       	=					model.generate(
									    config_and_inputs[1]['input_ids']    ,			num_beams=1    ,			max_length=3    ,			output_attentions=_snake_case    ,			return_dict_in_generate=_snake_case    ,			**_snake_case    ,			)
									# We check the state of decoder_attentions and cross_attentions just from the last step
									UpperCAmelCase_		:      Tuple       	=					out[attn_name] if attn_name == attention_names[0] else out[attn_name][-1]
									self.assertEqual(sum([w.sum().item() for w in attn_weights])    ,			0.0)





							@unittest.skip('Does not work on the tiny model as we keep hitting edge cases.')
							def 						_snake_case							(			self)							->				str:
								pass



@require_torch
@require_sentencepiece
@require_tokenizers
class 					lowercase						(  unittest.TestCase ):


							@slow
							@unittest.skip(
							    'Unless we stop stripping left and right by default for all special tokens, the expected ids obtained here will not match the original ones. Wait for https://github.com/huggingface/transformers/pull/23909 to be merged')
							def 						_snake_case							(			self)							->				Optional[int]:
								UpperCAmelCase_		:      Optional[int]       	=					UMTaForConditionalGeneration.from_pretrained('google/umt5-small'    ,			return_dict=_snake_case).to(_snake_case)
								UpperCAmelCase_		:      Optional[int]       	=					AutoTokenizer.from_pretrained('google/umt5-small'    ,			use_fast=_snake_case    ,			legacy=_snake_case)
								UpperCAmelCase_		:      Any       	=					[
								    'Bonjour monsieur <extra_id_0> bien <extra_id_1>.',
								    'No se como puedo <extra_id_0>.',
								    'This is the reason why we <extra_id_0> them.',
								    'The <extra_id_0> walks in <extra_id_1>, seats',
								    'A <extra_id_0> walks into a bar and orders a <extra_id_1> with <extra_id_2> pinch of <extra_id_3>.',
								]
								UpperCAmelCase_		:      List[Any]       	=					tokenizer(_snake_case    ,			return_tensors='pt'    ,			padding=_snake_case).input_ids
								# fmt: off
								UpperCAmelCase_		:      Optional[Any]       	=					torch.tensor(
								    [
								        [ 3_8530, 21_0703, 25_6299, 1410, 25_6298, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [   826, 321, 671, 2_5922, 25_6299, 274, 1, 0,0, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [  1460, 339, 312, 1_9014, 1_0620, 758, 25_6299, 2355,274, 1, 0, 0, 0, 0, 0, 0,0, 0],
								        [   517, 25_6299, 1_4869, 281, 301, 25_6298, 275, 11_9983,1, 0, 0, 0, 0, 0, 0, 0,0, 0],
								        [   320, 25_6299, 1_4869, 281, 2234, 289, 2275, 333,6_1391, 289, 25_6298, 543, 25_6297, 16_8714, 329, 25_6296,274, 1],
								    ])
								# fmt: on
								torch.testing.assert_allclose(_snake_case    ,			_snake_case)

								UpperCAmelCase_		:      Any       	=					model.generate(input_ids.to(_snake_case))
								UpperCAmelCase_		:      Any       	=					[
								    '<pad><extra_id_0> et<extra_id_1> [eod] <extra_id_2><extra_id_55>.. [eod] 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 💐 <extra_id_56>ajšietosto<extra_id_56>lleux<extra_id_19><extra_id_6>ajšie</s>',
								    '<pad><extra_id_0>.<extra_id_1>.,<0x0A>...spech <0x0A><extra_id_20> <extra_id_21></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>',
								    '<pad><extra_id_0> are not going to be a part of the world. We are not going to be a part of<extra_id_1> and<extra_id_2><0x0A><extra_id_48>.<extra_id_48></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>',
								    '<pad><extra_id_0> door<extra_id_1>, the door<extra_id_2> 피해[/</s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>',
								    '<pad><extra_id_0>nyone who<extra_id_1> drink<extra_id_2> a<extra_id_3> alcohol<extra_id_4> A<extra_id_5> A. This<extra_id_6> I<extra_id_7><extra_id_52><extra_id_53></s><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>',
								]
								UpperCAmelCase_		:      Optional[Any]       	=					tokenizer.batch_decode(_snake_case)
								self.assertEqual(_snake_case    ,			_snake_case)









'''simple docstring'''


import argparse
import csv
import logging
import os
import random

import numpy as np
import torch
from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
from tqdm import tqdm, trange

from transformers import (
    CONFIG_NAME,
    WEIGHTS_NAME,
    AdamW,
    OpenAIGPTDoubleHeadsModel,
    OpenAIGPTTokenizer,
    get_linear_schedule_with_warmup,
)


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







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase  )    ->    Union[str, Any]:
	UpperCAmelCase_		:      Any       	=					np.argmax(UpperCamelCase   ,axis=1  )
	return np.sum(outputs == labels  )







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase  )    ->    List[str]:
	with open(UpperCamelCase   ,encoding='utf_8'  ) as f:
		UpperCAmelCase_		:      int       	=					csv.reader(UpperCamelCase  )
		UpperCAmelCase_		:      Tuple       	=					[]
		next(UpperCamelCase  )  # skip the first line
		for line in tqdm(UpperCamelCase  ):
			output.append((' '.join(line[1:5]  ), line[5], line[6], int(line[-1]  ) - 1)  )
	return output







def 						SCREAMING_SNAKE_CASE(  UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase  )    ->    Optional[int]:
	UpperCAmelCase_		:      str       	=					[]
	for dataset in encoded_datasets:
		UpperCAmelCase_		:      str       	=					len(UpperCamelCase  )
		UpperCAmelCase_		:      str       	=					np.zeros((n_batch, 2, input_len)   ,dtype=np.intaa  )
		UpperCAmelCase_		:      List[Any]       	=					np.zeros((n_batch, 2)   ,dtype=np.intaa  )
		UpperCAmelCase_		:      Any       	=					np.full((n_batch, 2, input_len)   ,fill_value=-1_0_0   ,dtype=np.intaa  )
		UpperCAmelCase_		:      List[Any]       	=					np.zeros((n_batch,)   ,dtype=np.intaa  )
		for (
		    i,
		    (story, conta, conta, mc_label),
		) in enumerate(UpperCamelCase  ):
			UpperCAmelCase_		:      Union[str, Any]       	=					[start_token] + story[:cap_length] + [delimiter_token] + conta[:cap_length] + [clf_token]
			UpperCAmelCase_		:      List[Any]       	=					[start_token] + story[:cap_length] + [delimiter_token] + conta[:cap_length] + [clf_token]
			UpperCAmelCase_		:      Any       	=					with_conta
			UpperCAmelCase_		:      Tuple       	=					with_conta
			UpperCAmelCase_		:      Union[str, Any]       	=					len(UpperCamelCase  ) - 1
			UpperCAmelCase_		:      int       	=					len(UpperCamelCase  ) - 1
			UpperCAmelCase_		:      Any       	=					with_conta
			UpperCAmelCase_		:      Any       	=					with_conta
			UpperCAmelCase_		:      Optional[int]       	=					mc_label
		UpperCAmelCase_		:      Tuple       	=					(input_ids, mc_token_ids, lm_labels, mc_labels)
		tensor_datasets.append(tuple(torch.tensor(UpperCamelCase  ) for t in all_inputs  )  )
	return tensor_datasets







def 						SCREAMING_SNAKE_CASE(  )    ->    str:
	UpperCAmelCase_		:      Union[str, Any]       	=					argparse.ArgumentParser()
	parser.add_argument('--model_name'   ,type=UpperCamelCase   ,default='openai-gpt'   ,help='pretrained model name'  )
	parser.add_argument('--do_train'   ,action='store_true'   ,help='Whether to run training.'  )
	parser.add_argument('--do_eval'   ,action='store_true'   ,help='Whether to run eval on the dev set.'  )
	parser.add_argument(
	    '--output_dir'   ,default=UpperCamelCase   ,type=UpperCamelCase   ,required=UpperCamelCase   ,help='The output directory where the model predictions and checkpoints will be written.'   ,)
	parser.add_argument('--train_dataset'   ,type=UpperCamelCase   ,default=''  )
	parser.add_argument('--eval_dataset'   ,type=UpperCamelCase   ,default=''  )
	parser.add_argument('--seed'   ,type=UpperCamelCase   ,default=4_2  )
	parser.add_argument('--num_train_epochs'   ,type=UpperCamelCase   ,default=3  )
	parser.add_argument('--train_batch_size'   ,type=UpperCamelCase   ,default=8  )
	parser.add_argument('--eval_batch_size'   ,type=UpperCamelCase   ,default=1_6  )
	parser.add_argument('--adam_epsilon'   ,default=1e-8   ,type=UpperCamelCase   ,help='Epsilon for Adam optimizer.'  )
	parser.add_argument('--max_grad_norm'   ,type=UpperCamelCase   ,default=1  )
	parser.add_argument(
	    '--max_steps'   ,default=-1   ,type=UpperCamelCase   ,help=(
	        'If > 0: set total number of training                         steps to perform. Override num_train_epochs.'
	    )   ,)
	parser.add_argument(
	    '--gradient_accumulation_steps'   ,type=UpperCamelCase   ,default=1   ,help='Number of updates steps to accumulate before                        performing a backward/update pass.'   ,)
	parser.add_argument('--learning_rate'   ,type=UpperCamelCase   ,default=6.25e-5  )
	parser.add_argument('--warmup_steps'   ,default=0   ,type=UpperCamelCase   ,help='Linear warmup over warmup_steps.'  )
	parser.add_argument('--lr_schedule'   ,type=UpperCamelCase   ,default='warmup_linear'  )
	parser.add_argument('--weight_decay'   ,type=UpperCamelCase   ,default=0.01  )
	parser.add_argument('--lm_coef'   ,type=UpperCamelCase   ,default=0.9  )
	parser.add_argument('--n_valid'   ,type=UpperCamelCase   ,default=3_7_4  )

	parser.add_argument('--server_ip'   ,type=UpperCamelCase   ,default=''   ,help='Can be used for distant debugging.'  )
	parser.add_argument('--server_port'   ,type=UpperCamelCase   ,default=''   ,help='Can be used for distant debugging.'  )
	UpperCAmelCase_		:      Optional[Any]       	=					parser.parse_args()
	print(UpperCamelCase  )

	if args.server_ip and args.server_port:
		# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
		import ptvsd

		print('Waiting for debugger attach'  )
		ptvsd.enable_attach(address=(args.server_ip, args.server_port)   ,redirect_output=UpperCamelCase  )
		ptvsd.wait_for_attach()

	random.seed(args.seed  )
	np.random.seed(args.seed  )
	torch.manual_seed(args.seed  )
	torch.cuda.manual_seed_all(args.seed  )

	UpperCAmelCase_		:      Union[str, Any]       	=					torch.device('cuda' if torch.cuda.is_available() else 'cpu'  )
	UpperCAmelCase_		:      Dict       	=					torch.cuda.device_count()
	logger.info('device: {}, n_gpu {}'.format(UpperCamelCase   ,UpperCamelCase  )  )

	if not args.do_train and not args.do_eval:
		raise ValueError('At least one of `do_train` or `do_eval` must be True.'  )

	if not os.path.exists(args.output_dir  ):
		os.makedirs(args.output_dir  )

	# Load tokenizer and model
	# This loading functions also add new tokens and embeddings called `special tokens`
	# These new embeddings will be fine-tuned on the RocStories dataset
	UpperCAmelCase_		:      Optional[int]       	=					['_start_', '_delimiter_', '_classify_']
	UpperCAmelCase_		:      Union[str, Any]       	=					OpenAIGPTTokenizer.from_pretrained(args.model_name  )
	tokenizer.add_tokens(UpperCamelCase  )
	UpperCAmelCase_		:      Optional[Any]       	=					tokenizer.convert_tokens_to_ids(UpperCamelCase  )
	UpperCAmelCase_		:      Union[str, Any]       	=					OpenAIGPTDoubleHeadsModel.from_pretrained(args.model_name  )
	model.resize_token_embeddings(len(UpperCamelCase  )  )
	model.to(UpperCamelCase  )

	# Load and encode the datasets
	def tokenize_and_encode(UpperCamelCase  ):
		if isinstance(UpperCamelCase   ,UpperCamelCase  ):
			return tokenizer.convert_tokens_to_ids(tokenizer.tokenize(UpperCamelCase  )  )
		elif isinstance(UpperCamelCase   ,UpperCamelCase  ):
			return obj
		return [tokenize_and_encode(UpperCamelCase  ) for o in obj]

	logger.info('Encoding dataset...'  )
	UpperCAmelCase_		:      List[str]       	=					load_rocstories_dataset(args.train_dataset  )
	UpperCAmelCase_		:      Optional[Any]       	=					load_rocstories_dataset(args.eval_dataset  )
	UpperCAmelCase_		:      Dict       	=					(train_dataset, eval_dataset)
	UpperCAmelCase_		:      Union[str, Any]       	=					tokenize_and_encode(UpperCamelCase  )

	# Compute the max input length for the Transformer
	UpperCAmelCase_		:      Dict       	=					model.config.n_positions // 2 - 2
	UpperCAmelCase_		:      int       	=					max(
	    len(story[:max_length]  ) + max(len(conta[:max_length]  )   ,len(conta[:max_length]  )  ) + 3
	    for dataset in encoded_datasets
	    for story, conta, conta, _ in dataset  )
	UpperCAmelCase_		:      Dict       	=					min(UpperCamelCase   ,model.config.n_positions  )  # Max size of input for the pre-trained model

	# Prepare inputs tensors and dataloaders
	UpperCAmelCase_		:      Optional[int]       	=					pre_process_datasets(UpperCamelCase   ,UpperCamelCase   ,UpperCamelCase   ,*UpperCamelCase  )
	UpperCAmelCase_	,    UpperCAmelCase_		:      Dict       	=					tensor_datasets[0], tensor_datasets[1]

	UpperCAmelCase_		:      Optional[int]       	=					TensorDataset(*UpperCamelCase  )
	UpperCAmelCase_		:      Optional[Any]       	=					RandomSampler(UpperCamelCase  )
	UpperCAmelCase_		:      int       	=					DataLoader(UpperCamelCase   ,sampler=UpperCamelCase   ,batch_size=args.train_batch_size  )

	UpperCAmelCase_		:      int       	=					TensorDataset(*UpperCamelCase  )
	UpperCAmelCase_		:      Optional[int]       	=					SequentialSampler(UpperCamelCase  )
	UpperCAmelCase_		:      Tuple       	=					DataLoader(UpperCamelCase   ,sampler=UpperCamelCase   ,batch_size=args.eval_batch_size  )

	# Prepare optimizer
	if args.do_train:
		if args.max_steps > 0:
			UpperCAmelCase_		:      List[str]       	=					args.max_steps
			UpperCAmelCase_		:      List[str]       	=					args.max_steps // (len(UpperCamelCase  ) // args.gradient_accumulation_steps) + 1
		else:
			UpperCAmelCase_		:      List[str]       	=					len(UpperCamelCase  ) // args.gradient_accumulation_steps * args.num_train_epochs

		UpperCAmelCase_		:      Any       	=					list(model.named_parameters()  )
		UpperCAmelCase_		:      Any       	=					['bias', 'LayerNorm.bias', 'LayerNorm.weight']
		UpperCAmelCase_		:      Dict       	=					[
		    {
		        'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay  )],
		        'weight_decay': args.weight_decay,
		    },
		    {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay  )], 'weight_decay': 0.0},
		]
		UpperCAmelCase_		:      Any       	=					AdamW(UpperCamelCase   ,lr=args.learning_rate   ,eps=args.adam_epsilon  )
		UpperCAmelCase_		:      List[str]       	=					get_linear_schedule_with_warmup(
		    UpperCamelCase   ,num_warmup_steps=args.warmup_steps   ,num_training_steps=UpperCamelCase  )

	if args.do_train:
		UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_		:      List[Any]       	=					0, 0, None
		model.train()
		for _ in trange(int(args.num_train_epochs  )   ,desc='Epoch'  ):
			UpperCAmelCase_		:      Union[str, Any]       	=					0
			UpperCAmelCase_		:      Any       	=					0
			UpperCAmelCase_		:      List[Any]       	=					tqdm(UpperCamelCase   ,desc='Training'  )
			for step, batch in enumerate(UpperCamelCase  ):
				UpperCAmelCase_		:      Optional[Any]       	=					tuple(t.to(UpperCamelCase  ) for t in batch  )
				UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_		:      Tuple       	=					batch
				UpperCAmelCase_		:      Dict       	=					model(UpperCamelCase   ,mc_token_ids=UpperCamelCase   ,lm_labels=UpperCamelCase   ,mc_labels=UpperCamelCase  )
				UpperCAmelCase_		:      int       	=					args.lm_coef * losses[0] + losses[1]
				loss.backward()
				optimizer.step()
				scheduler.step()
				optimizer.zero_grad()
				tr_loss += loss.item()
				UpperCAmelCase_		:      Dict       	=					(
				    loss.item() if exp_average_loss is None else 0.7 * exp_average_loss + 0.3 * loss.item()
				)
				nb_tr_steps += 1
				UpperCAmelCase_		:      List[Any]       	=					'Training loss: {:.2e} lr: {:.2e}'.format(UpperCamelCase   ,scheduler.get_lr()[0]  )

    # Save a trained model
	if args.do_train:
		# Save a trained model, configuration and tokenizer
		UpperCAmelCase_		:      int       	=					model.module if hasattr(UpperCamelCase   ,'module'  ) else model  # Only save the model itself

		# If we save using the predefined names, we can load using `from_pretrained`
		UpperCAmelCase_		:      int       	=					os.path.join(args.output_dir   ,UpperCamelCase  )
		UpperCAmelCase_		:      Any       	=					os.path.join(args.output_dir   ,UpperCamelCase  )

		torch.save(model_to_save.state_dict()   ,UpperCamelCase  )
		model_to_save.config.to_json_file(UpperCamelCase  )
		tokenizer.save_vocabulary(args.output_dir  )

		# Load a trained model and vocabulary that you have fine-tuned
		UpperCAmelCase_		:      int       	=					OpenAIGPTDoubleHeadsModel.from_pretrained(args.output_dir  )
		UpperCAmelCase_		:      Dict       	=					OpenAIGPTTokenizer.from_pretrained(args.output_dir  )
		model.to(UpperCamelCase  )

	if args.do_eval:
		model.eval()
		UpperCAmelCase_	,    UpperCAmelCase_		:      Union[str, Any]       	=					0, 0
		UpperCAmelCase_	,    UpperCAmelCase_		:      Any       	=					0, 0
		for batch in tqdm(UpperCamelCase   ,desc='Evaluating'  ):
			UpperCAmelCase_		:      Union[str, Any]       	=					tuple(t.to(UpperCamelCase  ) for t in batch  )
			UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_		:      Any       	=					batch
			with torch.no_grad():
				UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_	,    UpperCAmelCase_		:      Union[str, Any]       	=					model(
				    UpperCamelCase   ,mc_token_ids=UpperCamelCase   ,lm_labels=UpperCamelCase   ,mc_labels=UpperCamelCase  )

			UpperCAmelCase_		:      Optional[Any]       	=					mc_logits.detach().cpu().numpy()
			UpperCAmelCase_		:      int       	=					mc_labels.to('cpu'  ).numpy()
			UpperCAmelCase_		:      Optional[int]       	=					accuracy(UpperCamelCase   ,UpperCamelCase  )

			eval_loss += mc_loss.mean().item()
			eval_accuracy += tmp_eval_accuracy

			nb_eval_examples += input_ids.size(0  )
			nb_eval_steps += 1

		UpperCAmelCase_		:      str       	=					eval_loss / nb_eval_steps
		UpperCAmelCase_		:      str       	=					eval_accuracy / nb_eval_examples
		UpperCAmelCase_		:      List[Any]       	=					tr_loss / nb_tr_steps if args.do_train else None
		UpperCAmelCase_		:      Dict       	=					{'eval_loss': eval_loss, 'eval_accuracy': eval_accuracy, 'train_loss': train_loss}

		UpperCAmelCase_		:      Union[str, Any]       	=					os.path.join(args.output_dir   ,'eval_results.txt'  )
		with open(UpperCamelCase   ,'w'  ) as writer:
			logger.info('***** Eval results *****'  )
			for key in sorted(result.keys()  ):
				logger.info('  %s = %s'   ,UpperCamelCase   ,str(result[key]  )  )
				writer.write('%s = %s\n' % (key, str(result[key]  ))  )


if __name__ == "__main__":
						main()






