

import unittest

from transformers import MraConfig, 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, floats_tensor, ids_tensor, random_attention_mask


if is_torch_available():
						import torch

						from transformers import (
						    MraForMaskedLM,
						    MraForMultipleChoice,
						    MraForQuestionAnswering,
						    MraForSequenceClassification,
						    MraForTokenClassification,
						    MraModel,
						)
						from transformers.models.mra.modeling_mra import MRA_PRETRAINED_MODEL_ARCHIVE_LIST



class   __lowercase    :




			"""simple docstring"""







			def __init__(			self  :					Optional[Any]       ,						lowerCAmelCase__  :					Tuple       ,						lowerCAmelCase__  :					Optional[int]=2       ,						lowerCAmelCase__  :					int=8       ,						lowerCAmelCase__  :					Tuple=True       ,						lowerCAmelCase__  :					Any=True       ,						lowerCAmelCase__  :					Optional[Any]=True       ,						lowerCAmelCase__  :					Dict=True       ,						lowerCAmelCase__  :					Union[str, Any]=99       ,						lowerCAmelCase__  :					Union[str, Any]=16       ,						lowerCAmelCase__  :					Dict=5       ,						lowerCAmelCase__  :					int=2       ,						lowerCAmelCase__  :					Union[str, Any]=36       ,						lowerCAmelCase__  :					str="gelu"       ,						lowerCAmelCase__  :					List[str]=0.0       ,						lowerCAmelCase__  :					Tuple=0.0       ,						lowerCAmelCase__  :					Any=512       ,						lowerCAmelCase__  :					Any=16       ,						lowerCAmelCase__  :					Optional[Any]=2       ,						lowerCAmelCase__  :					List[Any]=0.02       ,						lowerCAmelCase__  :					Tuple=3       ,						lowerCAmelCase__  :					List[str]=4       ,						lowerCAmelCase__  :					Optional[int]=None       ,						):
									SCREAMING_SNAKE_CASE_:			Dict  		=	parent
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	batch_size
									SCREAMING_SNAKE_CASE_:			Dict  		=	seq_length
									SCREAMING_SNAKE_CASE_:			str  		=	is_training
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	use_input_mask
									SCREAMING_SNAKE_CASE_:			int  		=	use_token_type_ids
									SCREAMING_SNAKE_CASE_:			int  		=	use_labels
									SCREAMING_SNAKE_CASE_:			Tuple  		=	vocab_size
									SCREAMING_SNAKE_CASE_:			Any  		=	hidden_size
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	num_hidden_layers
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	num_attention_heads
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	intermediate_size
									SCREAMING_SNAKE_CASE_:			str  		=	hidden_act
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	hidden_dropout_prob
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	max_position_embeddings
									SCREAMING_SNAKE_CASE_:			Tuple  		=	type_vocab_size
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	type_sequence_label_size
									SCREAMING_SNAKE_CASE_:			int  		=	initializer_range
									SCREAMING_SNAKE_CASE_:			Tuple  		=	num_labels
									SCREAMING_SNAKE_CASE_:			List[str]  		=	num_choices
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	scope



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	ids_tensor([self.batch_size, self.seq_length]       ,						self.vocab_size)

									SCREAMING_SNAKE_CASE_:			List[str]  		=	None
									if self.use_input_mask:
															SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	random_attention_mask([self.batch_size, self.seq_length])

									SCREAMING_SNAKE_CASE_:			Dict  		=	None
									if self.use_token_type_ids:
															SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	ids_tensor([self.batch_size, self.seq_length]       ,						self.type_vocab_size)

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

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	self.get_config()

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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					List[Any]):
									return MraConfig(
									    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       ,						is_decoder=UpperCAmelCase__       ,						initializer_range=self.initializer_range       ,						)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			str  		=	self.get_config()
									SCREAMING_SNAKE_CASE_:			str  		=	300
									return config



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									(
									    (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,     
									):			List[str]  		=	self.prepare_config_and_inputs()

									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	True
									SCREAMING_SNAKE_CASE_:			int  		=	floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	ids_tensor([self.batch_size, self.seq_length]       ,						vocab_size=2)

									return (
									    config,
									    input_ids,
									    token_type_ids,
									    input_mask,
									    sequence_labels,
									    token_labels,
									    choice_labels,
									    encoder_hidden_states,
									    encoder_attention_mask,
									)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					List[Any]       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					List[Any]       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					Dict):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	MraModel(config=UpperCAmelCase__)
									model.to(UpperCAmelCase__)
									model.eval()
									SCREAMING_SNAKE_CASE_:			Dict  		=	model(UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Any  		=	model(UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	model(UpperCAmelCase__)
									self.parent.assertEqual(result.last_hidden_state.shape       ,						(self.batch_size, self.seq_length, self.hidden_size))



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					List[Any]       ,						lowerCAmelCase__  :					Optional[int]       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					Optional[int]       ,						lowerCAmelCase__  :					Tuple       ,						lowerCAmelCase__  :					str       ,						):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	True
									SCREAMING_SNAKE_CASE_:			str  		=	MraModel(UpperCAmelCase__)
									model.to(UpperCAmelCase__)
									model.eval()
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	model(
									    UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__       ,						encoder_hidden_states=UpperCAmelCase__       ,						encoder_attention_mask=UpperCAmelCase__       ,						)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	model(
									    UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__       ,						encoder_hidden_states=UpperCAmelCase__       ,						)
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	model(UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__)
									self.parent.assertEqual(result.last_hidden_state.shape       ,						(self.batch_size, self.seq_length, self.hidden_size))



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					Any       ,						lowerCAmelCase__  :					Optional[Any]       ,						lowerCAmelCase__  :					Optional[Any]       ,						lowerCAmelCase__  :					str       ,						lowerCAmelCase__  :					Dict):
									SCREAMING_SNAKE_CASE_:			str  		=	MraForMaskedLM(config=UpperCAmelCase__)
									model.to(UpperCAmelCase__)
									model.eval()
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	model(UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__       ,						labels=UpperCAmelCase__)
									self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.vocab_size))



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



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]       ,						lowerCAmelCase__  :					Dict       ,						lowerCAmelCase__  :					int       ,						lowerCAmelCase__  :					Optional[int]       ,						lowerCAmelCase__  :					Tuple       ,						lowerCAmelCase__  :					Optional[Any]       ,						lowerCAmelCase__  :					List[str]       ,						lowerCAmelCase__  :					int):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	self.num_labels
									SCREAMING_SNAKE_CASE_:			Dict  		=	MraForTokenClassification(config=UpperCAmelCase__)
									model.to(UpperCAmelCase__)
									model.eval()
									SCREAMING_SNAKE_CASE_:			str  		=	model(UpperCAmelCase__       ,						attention_mask=UpperCAmelCase__       ,						token_type_ids=UpperCAmelCase__       ,						labels=UpperCAmelCase__)
									self.parent.assertEqual(result.logits.shape       ,						(self.batch_size, self.seq_length, self.num_labels))



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




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	self.prepare_config_and_inputs()
									(
									    (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,   (
									    SCREAMING_SNAKE_CASE_  
									)		,     
									):			Any  		=	config_and_inputs
									SCREAMING_SNAKE_CASE_:			str  		=	{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
									return config, inputs_dict







@require_torch
class   __lowercase    (			__lowercase		,	unittest.TestCase     ):




			"""simple docstring"""







			_UpperCAmelCase       :							str      							=					(
			    (
			        MraModel,
			        MraForMaskedLM,
			        MraForMultipleChoice,
			        MraForQuestionAnswering,
			        MraForSequenceClassification,
			        MraForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			_UpperCAmelCase       :							List[Any]      							=					False
			_UpperCAmelCase       :							Optional[int]      							=					False
			_UpperCAmelCase       :							List[str]      							=					False
			_UpperCAmelCase       :							Optional[int]      							=					False

			_UpperCAmelCase       :							str      							=					()



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			int  		=	MraModelTester(self)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	ConfigTester(self       ,						config_class=UpperCAmelCase__       ,						hidden_size=37)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					str):
									self.config_tester.run_common_tests()



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
															SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	type
															self.model_tester.create_and_check_model(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Any):
									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Union[str, Any]):
									SCREAMING_SNAKE_CASE_:			Dict  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_multiple_choice(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									SCREAMING_SNAKE_CASE_:			Any  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_question_answering(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									SCREAMING_SNAKE_CASE_:			Any  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_sequence_classification(*UpperCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					str):
									SCREAMING_SNAKE_CASE_:			Dict  		=	self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*UpperCAmelCase__)



			@slow
			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									for model_name in MRA_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															SCREAMING_SNAKE_CASE_:			Optional[int]  		=	MraModel.from_pretrained(UpperCAmelCase__)
															self.assertIsNotNone(UpperCAmelCase__)




			@unittest.skip(reason="MRA does not output attentions")
			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									return







@require_torch
class   __lowercase    (			unittest.TestCase     ):




			"""simple docstring"""







			@slow
			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]):
									SCREAMING_SNAKE_CASE_:			int  		=	MraModel.from_pretrained("uw-madison/mra-base-512-4")
									SCREAMING_SNAKE_CASE_:			str  		=	torch.arange(256).unsqueeze(0)

									with torch.no_grad():
															SCREAMING_SNAKE_CASE_:			Dict  		=	model(UpperCAmelCase__)[0]

									SCREAMING_SNAKE_CASE_:			Tuple  		=	torch.Size((1, 256, 768))
									self.assertEqual(output.shape       ,						UpperCAmelCase__)

									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	torch.tensor(
									    [[[-0.0140, 0.0830, -0.0381], [0.1546, 0.1402, 0.0220], [0.1162, 0.0851, 0.0165]]])

									self.assertTrue(torch.allclose(output[:, :3, :3]       ,						UpperCAmelCase__       ,						atol=1E-4))



			@slow
			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			int  		=	MraForMaskedLM.from_pretrained("uw-madison/mra-base-512-4")
									SCREAMING_SNAKE_CASE_:			List[str]  		=	torch.arange(256).unsqueeze(0)

									with torch.no_grad():
															SCREAMING_SNAKE_CASE_:			List[Any]  		=	model(UpperCAmelCase__)[0]

									SCREAMING_SNAKE_CASE_:			str  		=	5_0265

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	torch.Size((1, 256, vocab_size))
									self.assertEqual(output.shape       ,						UpperCAmelCase__)

									SCREAMING_SNAKE_CASE_:			List[Any]  		=	torch.tensor(
									    [[[9.2595, -3.6038, 11.8819], [9.3869, -3.2693, 11.0956], [11.8524, -3.4938, 13.1210]]])

									self.assertTrue(torch.allclose(output[:, :3, :3]       ,						UpperCAmelCase__       ,						atol=1E-4))




			@slow
			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									SCREAMING_SNAKE_CASE_:			List[str]  		=	MraForMaskedLM.from_pretrained("uw-madison/mra-base-4096-8-d3")
									SCREAMING_SNAKE_CASE_:			Dict  		=	torch.arange(4096).unsqueeze(0)

									with torch.no_grad():
															SCREAMING_SNAKE_CASE_:			int  		=	model(UpperCAmelCase__)[0]

									SCREAMING_SNAKE_CASE_:			List[Any]  		=	5_0265

									SCREAMING_SNAKE_CASE_:			Dict  		=	torch.Size((1, 4096, vocab_size))
									self.assertEqual(output.shape       ,						UpperCAmelCase__)

									SCREAMING_SNAKE_CASE_:			List[str]  		=	torch.tensor(
									    [[[5.4789, -2.3564, 7.5064], [7.9067, -1.3369, 9.9668], [9.0712, -1.8106, 7.0380]]])

									self.assertTrue(torch.allclose(output[:, :3, :3]       ,						UpperCAmelCase__       ,						atol=1E-4))









import unittest
from queue import Empty
from threading import Thread

from transformers import AutoTokenizer, TextIteratorStreamer, TextStreamer, is_torch_available
from transformers.testing_utils import CaptureStdout, require_torch, torch_device

from ..test_modeling_common import ids_tensor


if is_torch_available():
						import torch

						from transformers import AutoModelForCausalLM



@require_torch
class   __lowercase    (			unittest.TestCase     ):




			"""simple docstring"""







			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			str  		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
									SCREAMING_SNAKE_CASE_:			str  		=	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			str  		=	-1

									SCREAMING_SNAKE_CASE_:			Optional[int]  		=	ids_tensor((1, 5)       ,						vocab_size=model.config.vocab_size).to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Any  		=	model.generate(lowerCAmelCase__       ,						max_new_tokens=10       ,						do_sample=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[str]  		=	tokenizer.decode(greedy_ids[0])

									with CaptureStdout() as cs:
															SCREAMING_SNAKE_CASE_:			int  		=	TextStreamer(lowerCAmelCase__)
															model.generate(lowerCAmelCase__       ,						max_new_tokens=10       ,						do_sample=lowerCAmelCase__       ,						streamer=lowerCAmelCase__)
									# The greedy text should be printed to stdout, except for the final "\n" in the streamer
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	cs.out[:-1]

									self.assertEqual(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					int):
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
									SCREAMING_SNAKE_CASE_:			int  		=	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	-1

									SCREAMING_SNAKE_CASE_:			int  		=	ids_tensor((1, 5)       ,						vocab_size=model.config.vocab_size).to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			int  		=	model.generate(lowerCAmelCase__       ,						max_new_tokens=10       ,						do_sample=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	tokenizer.decode(greedy_ids[0])

									SCREAMING_SNAKE_CASE_:			int  		=	TextIteratorStreamer(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	{"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
									SCREAMING_SNAKE_CASE_:			Tuple  		=	Thread(target=model.generate       ,						kwargs=lowerCAmelCase__)
									thread.start()
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	""
									for new_text in streamer:
															streamer_text += new_text

									self.assertEqual(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[int]):
									SCREAMING_SNAKE_CASE_:			int  		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
									SCREAMING_SNAKE_CASE_:			int  		=	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	-1

									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	ids_tensor((1, 5)       ,						vocab_size=model.config.vocab_size).to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Dict  		=	model.generate(lowerCAmelCase__       ,						max_new_tokens=10       ,						do_sample=lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Tuple  		=	greedy_ids[:, input_ids.shape[1] :]
									SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	tokenizer.decode(new_greedy_ids[0])

									with CaptureStdout() as cs:
															SCREAMING_SNAKE_CASE_:			Dict  		=	TextStreamer(lowerCAmelCase__       ,						skip_prompt=lowerCAmelCase__)
															model.generate(lowerCAmelCase__       ,						max_new_tokens=10       ,						do_sample=lowerCAmelCase__       ,						streamer=lowerCAmelCase__)
									# The greedy text should be printed to stdout, except for the final "\n" in the streamer
									SCREAMING_SNAKE_CASE_:			Any  		=	cs.out[:-1]

									self.assertEqual(lowerCAmelCase__       ,						lowerCAmelCase__)



			def 	_SCREAMING_SNAKE_CASE      (			self  :					Tuple):
									# Tests that we can pass `decode_kwargs` to the streamer to control how the tokens are decoded. Must be tested
									# with actual models -- the dummy models' tokenizers are not aligned with their models, and
									# `skip_special_tokens=True` has no effect on them
									SCREAMING_SNAKE_CASE_:			Tuple  		=	AutoTokenizer.from_pretrained("distilgpt2")
									SCREAMING_SNAKE_CASE_:			List[str]  		=	AutoModelForCausalLM.from_pretrained("distilgpt2").to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	-1

									SCREAMING_SNAKE_CASE_:			List[str]  		=	torch.ones((1, 5)       ,						device=lowerCAmelCase__).long() * model.config.bos_token_id
									with CaptureStdout() as cs:
															SCREAMING_SNAKE_CASE_:			Union[str, Any]  		=	TextStreamer(lowerCAmelCase__       ,						skip_special_tokens=lowerCAmelCase__)
															model.generate(lowerCAmelCase__       ,						max_new_tokens=1       ,						do_sample=lowerCAmelCase__       ,						streamer=lowerCAmelCase__)

									# The prompt contains a special token, so the streamer should not print it. As such, the output text, when
									# re-tokenized, must only contain one token
									SCREAMING_SNAKE_CASE_:			str  		=	cs.out[:-1]  # Remove the final "\n"
									SCREAMING_SNAKE_CASE_:			Tuple  		=	tokenizer(lowerCAmelCase__       ,						return_tensors="pt")
									self.assertEqual(streamer_text_tokenized.input_ids.shape       ,						(1, 1))




			def 	_SCREAMING_SNAKE_CASE      (			self  :					Optional[Any]):
									SCREAMING_SNAKE_CASE_:			List[Any]  		=	AutoTokenizer.from_pretrained("hf-internal-testing/tiny-random-gpt2")
									SCREAMING_SNAKE_CASE_:			List[str]  		=	AutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-random-gpt2").to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Dict  		=	-1

									SCREAMING_SNAKE_CASE_:			List[str]  		=	ids_tensor((1, 5)       ,						vocab_size=model.config.vocab_size).to(lowerCAmelCase__)
									SCREAMING_SNAKE_CASE_:			Any  		=	TextIteratorStreamer(lowerCAmelCase__       ,						timeout=0.001)
									SCREAMING_SNAKE_CASE_:			Any  		=	{"input_ids": input_ids, "max_new_tokens": 10, "do_sample": False, "streamer": streamer}
									SCREAMING_SNAKE_CASE_:			Optional[Any]  		=	Thread(target=model.generate       ,						kwargs=lowerCAmelCase__)
									thread.start()

									# The streamer will timeout after 0.001 seconds, so an exception will be raised
									with self.assertRaises(lowerCAmelCase__):
															SCREAMING_SNAKE_CASE_:			Tuple  		=	""
															for new_text in streamer:
																					streamer_text += new_text






