


'''simple docstring'''



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  UpperCamelCase__(     unittest.TestCase ):




				def   a__(	self    :				Optional[int]   )->  Any:

										"""simple docstring"""




										UpperCAmelCase         =			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   )
										UpperCAmelCase         =			AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			-1

										UpperCAmelCase         =			ids_tensor((1, 5)   ,					vocab_size=model.config.vocab_size   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			model.generate(lowerCAmelCase   ,					max_new_tokens=10   ,					do_sample=lowerCAmelCase   )
										UpperCAmelCase         =			tokenizer.decode(greedy_ids[0]   )

										with CaptureStdout() as cs:
																UpperCAmelCase         =			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
										UpperCAmelCase         =			cs.out[:-1]

										self.assertEqual(lowerCAmelCase   ,					lowerCAmelCase   )




				def   a__(	self    :				Optional[int]   )->  str:

										"""simple docstring"""




										UpperCAmelCase         =			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   )
										UpperCAmelCase         =			AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			-1

										UpperCAmelCase         =			ids_tensor((1, 5)   ,					vocab_size=model.config.vocab_size   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			model.generate(lowerCAmelCase   ,					max_new_tokens=10   ,					do_sample=lowerCAmelCase   )
										UpperCAmelCase         =			tokenizer.decode(greedy_ids[0]   )

										UpperCAmelCase         =			TextIteratorStreamer(lowerCAmelCase   )
										UpperCAmelCase         =			{'''input_ids''': input_ids, '''max_new_tokens''': 10, '''do_sample''': False, '''streamer''': streamer}
										UpperCAmelCase         =			Thread(target=model.generate   ,					kwargs=lowerCAmelCase   )
										thread.start()
										UpperCAmelCase         =			''''''
										for new_text in streamer:
																streamer_text += new_text

										self.assertEqual(lowerCAmelCase   ,					lowerCAmelCase   )




				def   a__(	self    :				Tuple   )->  Union[str, Any]:

										"""simple docstring"""




										UpperCAmelCase         =			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   )
										UpperCAmelCase         =			AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			-1

										UpperCAmelCase         =			ids_tensor((1, 5)   ,					vocab_size=model.config.vocab_size   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			model.generate(lowerCAmelCase   ,					max_new_tokens=10   ,					do_sample=lowerCAmelCase   )
										UpperCAmelCase         =			greedy_ids[:, input_ids.shape[1] :]
										UpperCAmelCase         =			tokenizer.decode(new_greedy_ids[0]   )

										with CaptureStdout() as cs:
																UpperCAmelCase         =			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
										UpperCAmelCase         =			cs.out[:-1]

										self.assertEqual(lowerCAmelCase   ,					lowerCAmelCase   )




				def   a__(	self    :				Union[str, Any]   )->  List[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			AutoTokenizer.from_pretrained('''distilgpt2'''   )
										UpperCAmelCase         =			AutoModelForCausalLM.from_pretrained('''distilgpt2'''   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			-1

										UpperCAmelCase         =			torch.ones((1, 5)   ,					device=lowerCAmelCase   ).long() * model.config.bos_token_id
										with CaptureStdout() as cs:
																UpperCAmelCase         =			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
										UpperCAmelCase         =			cs.out[:-1]  # Remove the final "\n"
										UpperCAmelCase         =			tokenizer(lowerCAmelCase   ,					return_tensors='''pt'''   )
										self.assertEqual(streamer_text_tokenized.input_ids.shape   ,					(1, 1)   )





				def   a__(	self    :				Union[str, Any]   )->  Dict:

										"""simple docstring"""




										UpperCAmelCase         =			AutoTokenizer.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   )
										UpperCAmelCase         =			AutoModelForCausalLM.from_pretrained('''hf-internal-testing/tiny-random-gpt2'''   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			-1

										UpperCAmelCase         =			ids_tensor((1, 5)   ,					vocab_size=model.config.vocab_size   ).to(lowerCAmelCase   )
										UpperCAmelCase         =			TextIteratorStreamer(lowerCAmelCase   ,					timeout=0.001   )
										UpperCAmelCase         =			{'''input_ids''': input_ids, '''max_new_tokens''': 10, '''do_sample''': False, '''streamer''': streamer}
										UpperCAmelCase         =			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   ):
																UpperCAmelCase         =			''''''
																for new_text in streamer:
																						streamer_text += new_text









'''simple docstring'''



import time
import unittest

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

from ..test_modeling_common import ids_tensor


if is_torch_available():
	import torch

	from transformers.generation import (
	    MaxLengthCriteria,
	    MaxNewTokensCriteria,
	    MaxTimeCriteria,
	    StoppingCriteriaList,
	    validate_stopping_criteria,
	)






@require_torch
class  UpperCamelCase__(     unittest.TestCase ):




				def   a__(	self    :				List[Any]   ,					lowerCAmelCase    :				Optional[Any]   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase         =			3
										UpperCAmelCase         =			250

										UpperCAmelCase         =			ids_tensor((batch_size, length)   ,					lowerCAmelCase   )
										UpperCAmelCase         =			torch.ones((batch_size, length)   ,					device=lowerCAmelCase   ,					dtype=torch.float   ) / length
										return input_ids, scores




				def   a__(	self    :				Dict   )->  Optional[int]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			StoppingCriteriaList(
										    [
										        MaxLengthCriteria(max_length=10   ),
										        MaxTimeCriteria(max_time=0.1   ),
										    ]   )

										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




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

										"""simple docstring"""




										UpperCAmelCase         =			MaxLengthCriteria(max_length=10   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )




				def   a__(	self    :				Optional[Any]   )->  List[Any]:

										"""simple docstring"""




										UpperCAmelCase         =			MaxNewTokensCriteria(start_length=5   ,					max_new_tokens=5   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(9   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(10   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			StoppingCriteriaList([criteria]   )
										self.assertEqual(criteria_list.max_length   ,					10   )




				def   a__(	self    :				Tuple   )->  Optional[Any]:

										"""simple docstring"""




										UpperCAmelCase       ,       UpperCAmelCase         =			self._get_tensors(5   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   )
										self.assertFalse(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )

										UpperCAmelCase         =			MaxTimeCriteria(max_time=0.1   ,					initial_timestamp=time.time() - 0.2   )
										self.assertTrue(criteria(lowerCAmelCase   ,					lowerCAmelCase   )   )





				def   a__(	self    :				int   )->  Any:

										"""simple docstring"""




										validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					10   )

										with self.assertWarns(lowerCAmelCase   ):
																validate_stopping_criteria(StoppingCriteriaList([MaxLengthCriteria(10   )]   )   ,					11   )

										UpperCAmelCase         =			validate_stopping_criteria(StoppingCriteriaList()   ,					11   )

										self.assertEqual(len(lowerCAmelCase   )   ,					1   )





