






"""simple docstring"""
import math
import os
import unittest

from transformers import MegatronBertConfig, is_torch_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_sentencepiece, require_tokenizers, 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 (
			    MODEL_FOR_PRETRAINING_MAPPING,
			    MegatronBertForCausalLM,
			    MegatronBertForMaskedLM,
			    MegatronBertForMultipleChoice,
			    MegatronBertForNextSentencePrediction,
			    MegatronBertForPreTraining,
			    MegatronBertForQuestionAnswering,
			    MegatronBertForSequenceClassification,
			    MegatronBertForTokenClassification,
			    MegatronBertModel,
			)







class      UpperCamelCase_			:





	def __init__( self   :					List[Any]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Optional[Any]=13      ,      lowerCAmelCase_   :					int=7      ,      lowerCAmelCase_   :					List[Any]=True      ,      lowerCAmelCase_   :					Optional[Any]=True      ,      lowerCAmelCase_   :					Tuple=True      ,      lowerCAmelCase_   :					str=True      ,      lowerCAmelCase_   :					Optional[Any]=99      ,      lowerCAmelCase_   :					Optional[Any]=64      ,      lowerCAmelCase_   :					Optional[Any]=32      ,      lowerCAmelCase_   :					Union[str, Any]=5      ,      lowerCAmelCase_   :					Any=4      ,      lowerCAmelCase_   :					Dict=37      ,      lowerCAmelCase_   :					Optional[Any]="gelu"      ,      lowerCAmelCase_   :					Dict=0.1      ,      lowerCAmelCase_   :					Any=0.1      ,      lowerCAmelCase_   :					Tuple=512      ,      lowerCAmelCase_   :					Any=16      ,      lowerCAmelCase_   :					str=2      ,      lowerCAmelCase_   :					int=0.0_2      ,      lowerCAmelCase_   :					Dict=3      ,      lowerCAmelCase_   :					Union[str, Any]=4      ,      lowerCAmelCase_   :					str=None      ,      )		->     List[Any]:
				UpperCAmelCase_	:					Union[str, Any]     	=   parent
				UpperCAmelCase_	:					Optional[Any]     	=   batch_size
				UpperCAmelCase_	:					Tuple     	=   seq_length
				UpperCAmelCase_	:					Tuple     	=   is_training
				UpperCAmelCase_	:					Union[str, Any]     	=   use_input_mask
				UpperCAmelCase_	:					int     	=   use_token_type_ids
				UpperCAmelCase_	:					Tuple     	=   use_labels
				UpperCAmelCase_	:					Optional[Any]     	=   vocab_size
				UpperCAmelCase_	:					List[str]     	=   hidden_size
				UpperCAmelCase_	:					List[Any]     	=   embedding_size
				UpperCAmelCase_	:					List[str]     	=   num_hidden_layers
				UpperCAmelCase_	:					Optional[int]     	=   num_attention_heads
				UpperCAmelCase_	:					Optional[Any]     	=   intermediate_size
				UpperCAmelCase_	:					int     	=   hidden_act
				UpperCAmelCase_	:					Optional[Any]     	=   hidden_dropout_prob
				UpperCAmelCase_	:					List[str]     	=   attention_probs_dropout_prob
				UpperCAmelCase_	:					Optional[int]     	=   max_position_embeddings
				UpperCAmelCase_	:					Union[str, Any]     	=   type_vocab_size
				UpperCAmelCase_	:					List[str]     	=   type_sequence_label_size
				UpperCAmelCase_	:					List[str]     	=   initializer_range
				UpperCAmelCase_	:					Union[str, Any]     	=   num_labels
				UpperCAmelCase_	:					Any     	=   num_choices
				UpperCAmelCase_	:					Optional[int]     	=   scope





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]		)		->     Optional[int]:
				UpperCAmelCase_	:					List[str]     	=   ids_tensor([self.batch_size, self.seq_length]      ,      self.vocab_size		)

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

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

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

				UpperCAmelCase_	:					int     	=   self.get_config()

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     List[Any]:
				return MegatronBertConfig(
				    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      ,      embedding_size=self.embedding_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      ,      is_decoder=SCREAMING_SNAKE_CASE_      ,      initializer_range=self.initializer_range      ,      )





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					str		)		->     str:
				UpperCAmelCase_	:					Tuple     	=   MegatronBertModel(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					str     	=   model(SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_		)
				UpperCAmelCase_	:					Union[str, Any]     	=   model(SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_		)
				UpperCAmelCase_	:					Union[str, Any]     	=   model(SCREAMING_SNAKE_CASE_		)

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





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					int		)		->     Tuple:
				UpperCAmelCase_	:					Optional[int]     	=   MegatronBertForMaskedLM(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					Dict     	=   model(SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					str      ,      lowerCAmelCase_   :					Optional[int]		)		->     Optional[int]:
				UpperCAmelCase_	:					Any     	=   MegatronBertForCausalLM(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					Optional[Any]     	=   model(SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]		)		->     List[str]:
				UpperCAmelCase_	:					Union[str, Any]     	=   MegatronBertForNextSentencePrediction(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					List[str]     	=   model(
				    SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_      ,      )
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, 2)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					int      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Union[str, Any]		)		->     int:
				UpperCAmelCase_	:					int     	=   MegatronBertForPreTraining(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					List[str]     	=   model(
				    SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_      ,      next_sentence_label=SCREAMING_SNAKE_CASE_      ,      )
				self.parent.assertEqual(result.prediction_logits.shape      ,      (self.batch_size, self.seq_length, self.vocab_size)		)
				self.parent.assertEqual(result.seq_relationship_logits.shape      ,      (self.batch_size, 2)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					str      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					str		)		->     Tuple:
				UpperCAmelCase_	:					str     	=   MegatronBertForQuestionAnswering(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					Optional[Any]     	=   model(
				    SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      start_positions=SCREAMING_SNAKE_CASE_      ,      end_positions=SCREAMING_SNAKE_CASE_      ,      )
				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 	_SCREAMING_SNAKE_CASE ( self   :					Dict      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Dict      ,      lowerCAmelCase_   :					Tuple      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					int		)		->     List[str]:
				UpperCAmelCase_	:					Dict     	=   self.num_labels
				UpperCAmelCase_	:					Optional[int]     	=   MegatronBertForSequenceClassification(SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					Optional[Any]     	=   model(SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_labels)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					Any      ,      lowerCAmelCase_   :					Union[str, Any]		)		->     Optional[Any]:
				UpperCAmelCase_	:					List[Any]     	=   self.num_labels
				UpperCAmelCase_	:					Optional[int]     	=   MegatronBertForTokenClassification(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					Optional[Any]     	=   model(SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_		)
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.seq_length, self.num_labels)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]      ,      lowerCAmelCase_   :					Optional[int]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Union[str, Any]      ,      lowerCAmelCase_   :					List[str]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					Tuple		)		->     int:
				UpperCAmelCase_	:					Optional[Any]     	=   self.num_choices
				UpperCAmelCase_	:					Any     	=   MegatronBertForMultipleChoice(config=SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.eval()
				UpperCAmelCase_	:					str     	=   input_ids.unsqueeze(1		).expand(-1      ,      self.num_choices      ,      -1		).contiguous()
				UpperCAmelCase_	:					str     	=   token_type_ids.unsqueeze(1		).expand(-1      ,      self.num_choices      ,      -1		).contiguous()
				UpperCAmelCase_	:					int     	=   input_mask.unsqueeze(1		).expand(-1      ,      self.num_choices      ,      -1		).contiguous()
				UpperCAmelCase_	:					Dict     	=   model(
				    SCREAMING_SNAKE_CASE_      ,      attention_mask=SCREAMING_SNAKE_CASE_      ,      token_type_ids=SCREAMING_SNAKE_CASE_      ,      labels=SCREAMING_SNAKE_CASE_      ,      )
				self.parent.assertEqual(result.logits.shape      ,      (self.batch_size, self.num_choices)		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Any		)		->     Dict:
				UpperCAmelCase_	:					List[str]     	=   self.prepare_config_and_inputs()
				(
				    (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,      (
				    UpperCAmelCase_      
				)	,            
				)	:					Dict     	=   config_and_inputs
				UpperCAmelCase_	:					Tuple     	=   {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
				return config, inputs_dict







@require_torch
class      UpperCamelCase_			(__A							,		__A							,		unittest.TestCase	):
	__magic_name__       =							(
	    (
	        MegatronBertModel,
	        MegatronBertForMaskedLM,
	        MegatronBertForCausalLM,
	        MegatronBertForMultipleChoice,
	        MegatronBertForNextSentencePrediction,
	        MegatronBertForPreTraining,
	        MegatronBertForQuestionAnswering,
	        MegatronBertForSequenceClassification,
	        MegatronBertForTokenClassification,
	    )
	    if is_torch_available()
	    else ()
	)
	__magic_name__       =							(
	    {
	        'feature-extraction': MegatronBertModel,
	        'fill-mask': MegatronBertForMaskedLM,
	        'question-answering': MegatronBertForQuestionAnswering,
	        'text-classification': MegatronBertForSequenceClassification,
	        'text-generation': MegatronBertForCausalLM,
	        'token-classification': MegatronBertForTokenClassification,
	        'zero-shot': MegatronBertForSequenceClassification,
	    }
	    if is_torch_available()
	    else {}
	)
	__magic_name__       =							True
	# test_resize_embeddings = False
	__magic_name__       =							False





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[str]      ,      lowerCAmelCase_   :					List[Any]      ,      lowerCAmelCase_   :					Optional[Any]      ,      lowerCAmelCase_   :					List[Any]=False		)		->     int:
				UpperCAmelCase_	:					Dict     	=   super()._prepare_for_class(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      return_labels=SCREAMING_SNAKE_CASE_		)

				if return_labels:
							if model_class in get_values(SCREAMING_SNAKE_CASE_		):
										UpperCAmelCase_	:					List[Any]     	=   torch.zeros(
										    (self.model_tester.batch_size, self.model_tester.seq_length)      ,      dtype=torch.long      ,      device=SCREAMING_SNAKE_CASE_		)
										UpperCAmelCase_	:					Union[str, Any]     	=   torch.zeros(
										    self.model_tester.batch_size      ,      dtype=torch.long      ,      device=SCREAMING_SNAKE_CASE_		)
				return inputs_dict





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     List[Any]:
				UpperCAmelCase_	:					Optional[int]     	=   MegatronBertModelTester(self		)
				UpperCAmelCase_	:					Any     	=   ConfigTester(self      ,      config_class=SCREAMING_SNAKE_CASE_      ,      hidden_size=37		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Optional[int]:
				self.config_tester.run_common_tests()





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[Any]		)		->     Optional[Any]:
				UpperCAmelCase_	:					Tuple     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_model(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     Any:
				UpperCAmelCase_	:					int     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_masked_lm(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					int		)		->     int:
				UpperCAmelCase_	:					Any     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_multiple_choice(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Optional[int]		)		->     str:
				UpperCAmelCase_	:					Any     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_next_sequence_prediction(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Dict		)		->     Optional[int]:
				UpperCAmelCase_	:					Dict     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_pretraining(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     Optional[Any]:
				UpperCAmelCase_	:					List[str]     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_question_answering(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					List[Any]		)		->     Optional[int]:
				UpperCAmelCase_	:					Tuple     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_sequence_classification(*SCREAMING_SNAKE_CASE_		)





	def 	_SCREAMING_SNAKE_CASE ( self   :					Tuple		)		->     List[Any]:
				UpperCAmelCase_	:					Any     	=   self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_megatron_bert_for_token_classification(*SCREAMING_SNAKE_CASE_		)





def 							snake_case (  A__		):
			return torch.tensor(
			    snake_case__      ,dtype=torch.long      ,device=snake_case__      ,)


lowerCamelCase_							      = 1E-4







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





	@slow
	@unittest.skip("Model is not available."		)
	def 	_SCREAMING_SNAKE_CASE ( self   :					Union[str, Any]		)		->     Dict:
				UpperCAmelCase_	:					Tuple     	=   "nvidia/megatron-bert-uncased-345m"
				if "MYDIR" in os.environ:
							UpperCAmelCase_	:					Any     	=   os.path.join(os.environ["MYDIR"]      ,      SCREAMING_SNAKE_CASE_		)
				UpperCAmelCase_	:					Union[str, Any]     	=   MegatronBertModel.from_pretrained(SCREAMING_SNAKE_CASE_		)
				model.to(SCREAMING_SNAKE_CASE_		)
				model.half()
				UpperCAmelCase_	:					int     	=   _long_tensor([[101, 7_110, 1_005, 1_056, 2_023, 11_333, 17_413, 1_029, 102]]		)
				with torch.no_grad():
							UpperCAmelCase_	:					Tuple     	=   model(SCREAMING_SNAKE_CASE_		)[0]
				UpperCAmelCase_	:					Any     	=   torch.Size((1, 9, 1_024)		)
				self.assertEqual(output.shape      ,      SCREAMING_SNAKE_CASE_		)

				UpperCAmelCase_	:					int     	=   [-0.6_0_4_0, -0.2_5_1_7, -0.1_0_2_5, 0.3_4_2_0, -0.6_7_5_8, -0.0_0_1_7, -0.1_0_8_9, -0.1_9_9_0, 0.5_7_2_8]
				for ii in range(3		):
							for jj in range(3		):
										UpperCAmelCase_	:					Dict     	=   output[0, ii, jj]
										UpperCAmelCase_	:					Dict     	=   expected[3 * ii + jj]
										UpperCAmelCase_	:					int     	=   "ii={} jj={} a={} b={}".format(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_		)
										self.assertTrue(math.isclose(SCREAMING_SNAKE_CASE_      ,      SCREAMING_SNAKE_CASE_      ,      rel_tol=SCREAMING_SNAKE_CASE_      ,      abs_tol=SCREAMING_SNAKE_CASE_		)      ,      msg=SCREAMING_SNAKE_CASE_		)





import unittest

from transformers import (
    MODEL_FOR_CAUSAL_LM_MAPPING,
    TF_MODEL_FOR_CAUSAL_LM_MAPPING,
    TextGenerationPipeline,
    logging,
    pipeline,
)
from transformers.testing_utils import (
    CaptureLogger,
    is_pipeline_test,
    require_accelerate,
    require_tf,
    require_torch,
    require_torch_gpu,
    require_torch_or_tf,
)

from .test_pipelines_common import ANY
@is_pipeline_test
@require_torch_or_tf
class 					lowercase__(							unittest.TestCase				):



	"""simple docstring"""




	a     :List[Any]       	=      MODEL_FOR_CAUSAL_LM_MAPPING
	a     :Tuple       	=      TF_MODEL_FOR_CAUSAL_LM_MAPPING
	@require_torch
	def 		_lowercase    (			self    :     Tuple			)			->	Dict:
					lowercase_	    =   pipeline(task='''text-generation'''			,				model='''sshleifer/tiny-ctrl'''			,				framework='''pt'''			)
					# Using `do_sample=False` to force deterministic output
					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
					                ''' oscope. FiliFili@@'''
					            )
					        }
					    ]			,				)

					lowercase_	    =   text_generator(['''This is a test''', '''This is a second test''']			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {
					                '''generated_text''': (
					                    '''This is a test ☃ ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy oscope.'''
					                    ''' oscope. FiliFili@@'''
					                )
					            }
					        ],
					        [
					            {
					                '''generated_text''': (
					                    '''This is a second test ☃ segmental segmental segmental 议议eski eski flutter flutter Lacy'''
					                    ''' oscope. oscope. FiliFili@@'''
					                )
					            }
					        ],
					    ]			,				)

					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			,				num_return_sequences=2			,				return_tensors=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					    ]			,				)
					lowercase_	    =   text_generator.model.config.eos_token_id
					lowercase_	    =   '''<pad>'''
					lowercase_	    =   text_generator(
					    ['''This is a test''', '''This is a second test''']			,				do_sample=SCREAMING_SNAKE_CASE_			,				num_return_sequences=2			,				batch_size=2			,				return_tensors=SCREAMING_SNAKE_CASE_			,				)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        ],
					        [
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					            {'''generated_token_ids''': ANY(SCREAMING_SNAKE_CASE_			)},
					        ],
					    ]			,				)
	@require_tf
	def 		_lowercase    (			self    :     str			)			->	Optional[int]:
					lowercase_	    =   pipeline(task='''text-generation'''			,				model='''sshleifer/tiny-ctrl'''			,				framework='''tf'''			)

					# Using `do_sample=False` to force deterministic output
					lowercase_	    =   text_generator('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
					                ''' please,'''
					            )
					        }
					    ]			,				)

					lowercase_	    =   text_generator(['''This is a test''', '''This is a second test''']			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [
					            {
					                '''generated_text''': (
					                    '''This is a test FeyFeyFey(Croatis.), s.), Cannes Cannes Cannes 閲閲Cannes Cannes Cannes 攵'''
					                    ''' please,'''
					                )
					            }
					        ],
					        [
					            {
					                '''generated_text''': (
					                    '''This is a second test Chieftain Chieftain prefecture prefecture prefecture Cannes Cannes'''
					                    ''' Cannes 閲閲Cannes Cannes Cannes 攵 please,'''
					                )
					            }
					        ],
					    ]			,				)
	def 		_lowercase    (			self    :     Optional[Any]			,				SCREAMING_SNAKE_CASE_    :     int			,				SCREAMING_SNAKE_CASE_    :     Union[str, Any]			,				SCREAMING_SNAKE_CASE_    :     List[str]			)			->	Union[str, Any]:
					lowercase_	    =   TextGenerationPipeline(model=SCREAMING_SNAKE_CASE_			,				tokenizer=SCREAMING_SNAKE_CASE_			)
					return text_generator, ["This is a test", "Another test"]
	def 		_lowercase    (			self    :     List[Any]			)			->	Optional[int]:
					lowercase_	    =   '''Hello I believe in'''
					lowercase_	    =   pipeline('''text-generation'''			,				model='''hf-internal-testing/tiny-random-gpt2'''			)
					lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': '''Hello I believe in fe fe fe fe fe fe fe fe fe fe fe fe'''}]			,				)

					lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				stop_sequence=''' fe'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': '''Hello I believe in fe'''}]			)
	def 		_lowercase    (			self    :     List[Any]			,				SCREAMING_SNAKE_CASE_    :     Tuple			,				SCREAMING_SNAKE_CASE_    :     List[str]			)			->	str:
					lowercase_	    =   text_generator.model
					lowercase_	    =   text_generator.tokenizer

					lowercase_	    =   text_generator('''This is a test'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test'''			)			)

					lowercase_	    =   text_generator('''This is a test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertNotIn('''This is a test'''			,				outputs[0]['''generated_text''']			)

					lowercase_	    =   pipeline(task='''text-generation'''			,				model=SCREAMING_SNAKE_CASE_			,				tokenizer=SCREAMING_SNAKE_CASE_			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					lowercase_	    =   text_generator('''This is a test'''			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertNotIn('''This is a test'''			,				outputs[0]['''generated_text''']			)

					lowercase_	    =   text_generator('''This is a test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					self.assertTrue(outputs[0]['''generated_text'''].startswith('''This is a test'''			)			)

					lowercase_	    =   text_generator(['''This is great !''', '''Something else''']			,				num_return_sequences=2			,				do_sample=SCREAMING_SNAKE_CASE_			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
					        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
					    ]			,				)

					if text_generator.tokenizer.pad_token is not None:
									lowercase_	    =   text_generator(
									    ['''This is great !''', '''Something else''']			,				num_return_sequences=2			,				batch_size=2			,				do_sample=SCREAMING_SNAKE_CASE_			)
									self.assertEqual(
									    SCREAMING_SNAKE_CASE_			,				[
									        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
									        [{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}, {'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}],
									    ]			,				)

					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			,				return_text=SCREAMING_SNAKE_CASE_			)
					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_full_text=SCREAMING_SNAKE_CASE_			,				return_tensors=SCREAMING_SNAKE_CASE_			)
					with self.assertRaises(SCREAMING_SNAKE_CASE_			):
									lowercase_	    =   text_generator('''test'''			,				return_text=SCREAMING_SNAKE_CASE_			,				return_tensors=SCREAMING_SNAKE_CASE_			)

					# Empty prompt is slighly special
					# it requires BOS token to exist.
					# Special case for Pegasus which will always append EOS so will
					# work even without BOS.
					if (
					    text_generator.tokenizer.bos_token_id is not None
					    or "Pegasus" in tokenizer.__class__.__name__
					    or "Git" in model.__class__.__name__
					):
									lowercase_	    =   text_generator(''''''			)
									self.assertEqual(SCREAMING_SNAKE_CASE_			,				[{'''generated_text''': ANY(SCREAMING_SNAKE_CASE_			)}]			)
					else:
									with self.assertRaises((ValueError, AssertionError)			):
													lowercase_	    =   text_generator(''''''			)

					if text_generator.framework == "tf":
									# TF generation does not support max_new_tokens, and it's impossible
									# to control long generation with only max_length without
									# fancy calculation, dismissing tests for now.
									return
					# We don't care about infinite range models.
					# They already work.
					# Skip this test for XGLM, since it uses sinusoidal positional embeddings which are resized on-the-fly.
					lowercase_	    =   ['''RwkvForCausalLM''', '''XGLMForCausalLM''', '''GPTNeoXForCausalLM''']
					if (
					    tokenizer.model_max_length < 1_0_0_0_0
					    and text_generator.model.__class__.__name__ not in EXTRA_MODELS_CAN_HANDLE_LONG_INPUTS
					):
									# Handling of large generations
									with self.assertRaises((RuntimeError, IndexError, ValueError, AssertionError)			):
													text_generator('''This is a test''' * 5_0_0			,				max_new_tokens=2_0			)

									lowercase_	    =   text_generator('''This is a test''' * 5_0_0			,				handle_long_generation='''hole'''			,				max_new_tokens=2_0			)
									# Hole strategy cannot work
									with self.assertRaises(SCREAMING_SNAKE_CASE_			):
													text_generator(
													    '''This is a test''' * 5_0_0			,				handle_long_generation='''hole'''			,				max_new_tokens=tokenizer.model_max_length + 1_0			,				)
	@require_torch
	@require_accelerate
	@require_torch_gpu
	def 		_lowercase    (			self    :     Union[str, Any]			)			->	Tuple:
					import torch

					# Classic `model_kwargs`
					lowercase_	    =   pipeline(
					    model='''hf-internal-testing/tiny-random-bloom'''			,				model_kwargs={'''device_map''': '''auto''', '''torch_dtype''': torch.bfloataa}			,				)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.bfloataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)

					# Upgraded those two to real pipeline arguments (they just get sent for the model as they're unlikely to mean anything else.)
					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			,				torch_dtype=torch.bfloataa			)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.bfloataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)

					# torch_dtype will be automatically set to float32 if not provided - check: https://github.com/huggingface/transformers/pull/20602
					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			)
					self.assertEqual(pipe.model.device			,				torch.device(0			)			)
					self.assertEqual(pipe.model.lm_head.weight.dtype			,				torch.floataa			)
					lowercase_	    =   pipe('''This is a test'''			)
					self.assertEqual(
					    SCREAMING_SNAKE_CASE_			,				[
					        {
					            '''generated_text''': (
					                '''This is a test test test test test test test test test test test test test test test test'''
					                ''' test'''
					            )
					        }
					    ]			,				)
	@require_torch
	@require_torch_gpu
	def 		_lowercase    (			self    :     Optional[int]			)			->	List[Any]:
					import torch

					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device=0			,				torch_dtype=torch.floataa			)
					pipe('''This is a test'''			)
	@require_torch
	@require_accelerate
	@require_torch_gpu
	def 		_lowercase    (			self    :     int			)			->	Dict:
					import torch

					lowercase_	    =   pipeline(model='''hf-internal-testing/tiny-random-bloom'''			,				device_map='''auto'''			,				torch_dtype=torch.floataa			)
					pipe('''This is a test'''			,				do_sample=SCREAMING_SNAKE_CASE_			,				top_p=0.5			)




	def 		_lowercase    (			self    :     List[Any]			)			->	int:
					lowercase_	    =   '''Hello world'''
					lowercase_	    =   pipeline('''text-generation'''			,				model='''hf-internal-testing/tiny-random-gpt2'''			)
					if text_generator.model.framework == "tf":
									lowercase_	    =   logging.get_logger('''transformers.generation.tf_utils'''			)
					else:
									lowercase_	    =   logging.get_logger('''transformers.generation.utils'''			)
					lowercase_	    =   '''Both `max_new_tokens`'''  # The beggining of the message to be checked in this test

					# Both are set by the user -> log warning
					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_length=1_0			,				max_new_tokens=1			)
					self.assertIn(SCREAMING_SNAKE_CASE_			,				cl.out			)

					# The user only sets one -> no warning
					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_new_tokens=1			)
					self.assertNotIn(SCREAMING_SNAKE_CASE_			,				cl.out			)

					with CaptureLogger(SCREAMING_SNAKE_CASE_			) as cl:
									lowercase_	    =   text_generator(SCREAMING_SNAKE_CASE_			,				max_length=1_0			)
					self.assertNotIn(SCREAMING_SNAKE_CASE_			,				cl.out			)




