
import unittest

from transformers import ErnieMConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch

					from transformers import (
					    ErnieMForInformationExtraction,
					    ErnieMForMultipleChoice,
					    ErnieMForQuestionAnswering,
					    ErnieMForSequenceClassification,
					    ErnieMForTokenClassification,
					    ErnieMModel,
					)
					from transformers.models.ernie_m.modeling_ernie_m import ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST


class 							lowercase   :



		def __init__(					self							, A_							, A_=13							, A_=7							, A_=True							, A_=True							, A_=True							, A_=99							, A_=32							, A_=5							, A_=4							, A_=37							, A_="gelu"							, A_=0.1							, A_=0.1							, A_=512							, A_=16							, A_=2							, A_=0.02							, A_=3							, A_=4							, A_=None							, )    ->   List[Any]:
				"""simple docstring"""
				UpperCamelCase            =	parent
				UpperCamelCase            =	batch_size
				UpperCamelCase            =	seq_length
				UpperCamelCase            =	is_training
				UpperCamelCase            =	use_input_mask
				UpperCamelCase            =	use_labels
				UpperCamelCase            =	vocab_size
				UpperCamelCase            =	hidden_size
				UpperCamelCase            =	num_hidden_layers
				UpperCamelCase            =	num_attention_heads
				UpperCamelCase            =	intermediate_size
				UpperCamelCase            =	hidden_act
				UpperCamelCase            =	hidden_dropout_prob
				UpperCamelCase            =	attention_probs_dropout_prob
				UpperCamelCase            =	max_position_embeddings
				UpperCamelCase            =	type_vocab_size
				UpperCamelCase            =	type_sequence_label_size
				UpperCamelCase            =	initializer_range
				UpperCamelCase            =	num_labels
				UpperCamelCase            =	num_choices
				UpperCamelCase            =	scope



		def __UpperCamelCase					(					self			)    ->   Dict:
				"""simple docstring"""
				UpperCamelCase            =	ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size			)

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

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

				UpperCamelCase            =	self.get_config()

				return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels



		def __UpperCamelCase					(					self			)    ->   int:
				"""simple docstring"""
				UpperCamelCase            =	ids_tensor([self.batch_size, self.seq_length]							, self.vocab_size			)

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

				return config, input_ids, input_mask



		def __UpperCamelCase					(					self			)    ->   Optional[Any]:
				"""simple docstring"""
				return ErnieMConfig(
				    vocab_size=self.vocab_size							, hidden_size=self.hidden_size							, num_hidden_layers=self.num_hidden_layers							, num_attention_heads=self.num_attention_heads							, intermediate_size=self.intermediate_size							, hidden_act=self.hidden_act							, hidden_dropout_prob=self.hidden_dropout_prob							, attention_probs_dropout_prob=self.attention_probs_dropout_prob							, max_position_embeddings=self.max_position_embeddings							, type_vocab_size=self.type_vocab_size							, initializer_range=self.initializer_range							, )



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   List[Any]:
				"""simple docstring"""
				UpperCamelCase            =	ErnieMModel(config=A_			)
				model.to(A_			)
				model.eval()
				UpperCamelCase            =	model(A_							, return_dict=A_			)
				self.parent.assertEqual(result.last_hidden_state.shape							, (self.batch_size, self.seq_length, self.hidden_size)			)



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase            =	ErnieMForQuestionAnswering(config=A_			)
				model.to(A_			)
				model.eval()
				UpperCamelCase            =	model(
				    A_							, attention_mask=A_							, start_positions=A_							, end_positions=A_							, )
				self.parent.assertEqual(result.start_logits.shape							, (self.batch_size, self.seq_length)			)
				self.parent.assertEqual(result.end_logits.shape							, (self.batch_size, self.seq_length)			)



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   Tuple:
				"""simple docstring"""
				UpperCamelCase            =	ErnieMForInformationExtraction(config=A_			)
				model.to(A_			)
				model.eval()
				UpperCamelCase            =	torch.ones_like(A_							, dtype=torch.floataa			)
				UpperCamelCase            =	model(
				    A_							, attention_mask=A_							, start_positions=A_							, end_positions=A_							, )
				self.parent.assertEqual(result.start_logits.shape							, (self.batch_size, self.seq_length)			)
				self.parent.assertEqual(result.end_logits.shape							, (self.batch_size, self.seq_length)			)



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   int:
				"""simple docstring"""
				UpperCamelCase            =	self.num_labels
				UpperCamelCase            =	ErnieMForSequenceClassification(A_			)
				model.to(A_			)
				model.eval()
				UpperCamelCase            =	model(A_							, attention_mask=A_							, labels=A_			)
				self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_labels)			)



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   Dict:
				"""simple docstring"""
				UpperCamelCase            =	self.num_labels
				UpperCamelCase            =	ErnieMForTokenClassification(config=A_			)
				model.to(A_			)
				model.eval()
				input_ids.to(A_			)
				input_mask.to(A_			)
				token_labels.to(A_			)

				UpperCamelCase            =	model(A_							, attention_mask=A_							, labels=A_			)

				self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.seq_length, self.num_labels)			)



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_							, A_			)    ->   Tuple:
				"""simple docstring"""
				UpperCamelCase            =	self.num_choices
				UpperCamelCase            =	ErnieMForMultipleChoice(config=A_			)
				model.to(A_			)
				model.eval()
				UpperCamelCase            =	input_ids.unsqueeze(1			).expand(-1							, self.num_choices							, -1			).contiguous()
				UpperCamelCase            =	input_mask.unsqueeze(1			).expand(-1							, self.num_choices							, -1			).contiguous()
				UpperCamelCase            =	model(
				    A_							, attention_mask=A_							, labels=A_							, )
				self.parent.assertEqual(result.logits.shape							, (self.batch_size, self.num_choices)			)



		def __UpperCamelCase					(					self			)    ->   Optional[int]:
				"""simple docstring"""
				UpperCamelCase            =	self.prepare_config_and_inputs()
				(
				    (
				    UpperCamelCase
				)      ,    (
				    UpperCamelCase
				)      ,    (
				    UpperCamelCase
				)      ,    (
				    UpperCamelCase
				)      ,    (
				    UpperCamelCase
				)      ,    (
				    UpperCamelCase
				)      ,    
				)            =	config_and_inputs
				UpperCamelCase            =	{'input_ids': input_ids, 'attention_mask': input_mask}
				return config, inputs_dict






@require_torch
class 							lowercase   (						_SCREAMING_SNAKE_CASE    ,  _SCREAMING_SNAKE_CASE    ,  unittest.TestCase  ):
		__lowercase			:			int      	=							(
		    (
		        ErnieMModel,
		        ErnieMForMultipleChoice,
		        ErnieMForQuestionAnswering,
		        ErnieMForSequenceClassification,
		        ErnieMForTokenClassification,
		    )
		    if is_torch_available()
		    else ()
		)
		__lowercase			:			Any      	=							()
		__lowercase			:			List[Any]      	=							(
		    {
		        "feature-extraction": ErnieMModel,
		        "question-answering": ErnieMForQuestionAnswering,
		        "text-classification": ErnieMForSequenceClassification,
		        "token-classification": ErnieMForTokenClassification,
		        "zero-shot": ErnieMForSequenceClassification,
		    }
		    if is_torch_available()
		    else {}
		)
		__lowercase			:			str      	=							False



		def __UpperCamelCase					(					self							, A_							, A_							, A_							, A_							, A_			)    ->   int:
				"""simple docstring"""
				if pipeline_test_casse_name == "QAPipelineTests":
						return True

				return False



		def __UpperCamelCase					(					self			)    ->   Tuple:
				"""simple docstring"""
				UpperCamelCase            =	ErnieMModelTester(self			)
				UpperCamelCase            =	ConfigTester(self							, config_class=A_							, hidden_size=37			)



		def __UpperCamelCase					(					self			)    ->   List[Any]:
				"""simple docstring"""
				self.config_tester.run_common_tests()



		def __UpperCamelCase					(					self			)    ->   Optional[int]:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*A_			)



		def __UpperCamelCase					(					self			)    ->   str:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				for type in ["absolute", "relative_key", "relative_key_query"]:
						UpperCamelCase            =	type
						self.model_tester.create_and_check_model(*A_			)



		def __UpperCamelCase					(					self			)    ->   Any:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_multiple_choice(*A_			)



		def __UpperCamelCase					(					self			)    ->   Dict:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_question_answering(*A_			)



		def __UpperCamelCase					(					self			)    ->   Any:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_sequence_classification(*A_			)



		def __UpperCamelCase					(					self			)    ->   str:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_information_extraction(*A_			)



		def __UpperCamelCase					(					self			)    ->   Union[str, Any]:
				"""simple docstring"""
				UpperCamelCase            =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*A_			)



		@slow
		def __UpperCamelCase					(					self			)    ->   Any:
				"""simple docstring"""
				for model_name in ERNIE_M_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						UpperCamelCase            =	ErnieMModel.from_pretrained(A_			)
						self.assertIsNotNone(A_			)






@require_torch
class 							lowercase   (						unittest.TestCase  ):



		@slow
		def __UpperCamelCase					(					self			)    ->   Optional[Any]:
				"""simple docstring"""
				UpperCamelCase            =	ErnieMModel.from_pretrained('susnato/ernie-m-base_pytorch'			)
				model.eval()
				UpperCamelCase            =	torch.tensor([[0, 1, 2, 3, 4, 5]]			)
				UpperCamelCase            =	model(A_			)[0]

				# TODO Replace vocab size
				UpperCamelCase            =	768

				UpperCamelCase            =	torch.Size((1, 6, hidden_size)			)
				self.assertEqual(output.shape							, A_			)

				UpperCamelCase            =	torch.tensor(
				    [[[-0.0012, 0.1245, -0.0214], [-0.0742, 0.0244, -0.0771], [-0.0333, 0.1164, -0.1554]]]			)

				self.assertTrue(torch.allclose(output[:, :3, :3]							, A_							, atol=1e-3			)			)


import argparse
import logging
import math
import time

import torch

from transformers import TransfoXLCorpus, TransfoXLLMHeadModel


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


def  A			(   )  ->    List[Any]:
		'''simple docstring'''

		UpperCamelCase            =	argparse.ArgumentParser(description='PyTorch Transformer Language Model'					)
		parser.add_argument('--model_name'       ,			type=lowercase       ,			default='transfo-xl-wt103'       ,			help='pretrained model name'					)
		parser.add_argument(
		    '--split'       ,			type=lowercase       ,			default='test'       ,			choices=['all', 'valid', 'test']       ,			help='which split to evaluate'					)
		parser.add_argument('--batch_size'       ,			type=lowercase       ,			default=10       ,			help='batch size'					)
		parser.add_argument('--tgt_len'       ,			type=lowercase       ,			default=128       ,			help='number of tokens to predict'					)
		parser.add_argument('--ext_len'       ,			type=lowercase       ,			default=0       ,			help='length of the extended context'					)
		parser.add_argument('--mem_len'       ,			type=lowercase       ,			default=1_600       ,			help='length of the retained previous heads'					)
		parser.add_argument('--clamp_len'       ,			type=lowercase       ,			default=1_000       ,			help='max positional embedding index'					)
		parser.add_argument('--no_cuda'       ,			action='store_true'       ,			help='Do not use CUDA even though CUA is available'					)
		parser.add_argument('--work_dir'       ,			type=lowercase       ,			required=lowercase       ,			help='path to the work_dir'					)
		parser.add_argument('--no_log'       ,			action='store_true'       ,			help='do not log the eval result'					)
		parser.add_argument('--same_length'       ,			action='store_true'       ,			help='set same length attention with masking'					)
		parser.add_argument('--server_ip'       ,			type=lowercase       ,			default=''       ,			help='Can be used for distant debugging.'					)
		parser.add_argument('--server_port'       ,			type=lowercase       ,			default=''       ,			help='Can be used for distant debugging.'					)
		UpperCamelCase            =	parser.parse_args()
		assert args.ext_len >= 0, "extended context length must be non-negative"

		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=lowercase					)
				ptvsd.wait_for_attach()

		UpperCamelCase            =	torch.device('cuda' if torch.cuda.is_available() and not args.no_cuda else 'cpu'					)
		logger.info('device: {}'.format(lowercase					)					)

		# Load a pre-processed dataset
		# You can also build the corpus yourself using TransfoXLCorpus methods
		# The pre-processing involve computing word frequencies to prepare the Adaptive input and SoftMax
		# and tokenizing the dataset
		# The pre-processed corpus is a convertion (using the conversion script )
		UpperCamelCase            =	TransfoXLCorpus.from_pretrained(args.model_name					)

		UpperCamelCase            =	corpus.get_iterator('valid'       ,			args.batch_size       ,			args.tgt_len       ,			device=lowercase       ,			ext_len=args.ext_len					)
		UpperCamelCase            =	corpus.get_iterator('test'       ,			args.batch_size       ,			args.tgt_len       ,			device=lowercase       ,			ext_len=args.ext_len					)

		# Load a pre-trained model
		UpperCamelCase            =	TransfoXLLMHeadModel.from_pretrained(args.model_name					)
		model.to(lowercase					)

		logger.info(
		    'Evaluating with bsz {} tgt_len {} ext_len {} mem_len {} clamp_len {}'.format(
		        args.batch_size       ,			args.tgt_len       ,			args.ext_len       ,			args.mem_len       ,			args.clamp_len					)					)

		model.reset_memory_length(args.mem_len					)
		if args.clamp_len > 0:
				UpperCamelCase            =	args.clamp_len
		if args.same_length:
				UpperCamelCase            =	True

		###############################################################################
		# Evaluation code
		###############################################################################
		def evaluate(lowercase					):
				# Turn on evaluation mode which disables dropout.
				model.eval()
				UpperCamelCase      ,    UpperCamelCase            =	0, 0.0
				UpperCamelCase            =	time.time()
				with torch.no_grad():
						UpperCamelCase            =	None
						for idx, (data, target, seq_len) in enumerate(lowercase					):
								UpperCamelCase            =	model(lowercase       ,			lm_labels=lowercase       ,			mems=lowercase					)
								UpperCamelCase      ,    UpperCamelCase      ,    UpperCamelCase            =	ret
								UpperCamelCase            =	loss.mean()
								total_loss += seq_len * loss.item()
								total_len += seq_len
						UpperCamelCase            =	time.time() - start_time
				logger.info('Time : {:.2f}s, {:.2f}ms/segment'.format(lowercase       ,			1_000 * total_time / (idx + 1)					)					)
				return total_loss / total_len

		# Run on test data.
		if args.split == "all":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	evaluate(lowercase					)
		elif args.split == "valid":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	None
		elif args.split == "test":
				UpperCamelCase            =	evaluate(lowercase					)
				UpperCamelCase            =	None

		def format_log(lowercase       ,			lowercase					):
				UpperCamelCase            =	'| {0} loss {1:5.2f} | {0} ppl {2:9.3f} '.format(lowercase       ,			lowercase       ,			math.exp(lowercase					)					)
				return log_str

		UpperCamelCase            =	''
		if valid_loss is not None:
				log_str += format_log(lowercase       ,			'valid'					)
		if test_loss is not None:
				log_str += format_log(lowercase       ,			'test'					)

		logger.info('=' * 100					)
		logger.info(lowercase					)
		logger.info('=' * 100					)


if __name__ == "__main__":
					main()
