
'''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     __snake_case		(   unittest.TestCase):
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Any									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			Optional[int]									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Union[str, Any]									= -1

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

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

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			List[str]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			Dict									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			List[Any]									= -1

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

									lowerCamelCase					:			Union[str, Any]									= TextIteratorStreamer(A      )
									lowerCamelCase					:			Tuple									= {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
									lowerCamelCase					:			List[Any]									= Thread(target=model.generate,    kwargs=A      )
									thread.start()
									lowerCamelCase					:			List[Any]									= ''
									for new_text in streamer:
													streamer_text += new_text

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			int									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Optional[Any]									= -1

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

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

									self.assertEqual(A,    A      )
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[Any]									= AutoTokenizer.from_pretrained('distilgpt2'      )
									lowerCamelCase					:			Dict									= AutoModelForCausalLM.from_pretrained('distilgpt2'      ).to(A      )
									lowerCamelCase					:			Dict									= -1

									lowerCamelCase					:			Any									= torch.ones((1, 5),    device=A      ).long() * model.config.bos_token_id
									with CaptureStdout() as cs:
													lowerCamelCase					:			int									= TextStreamer(A,    skip_special_tokens=A      )
													model.generate(A,    max_new_tokens=1,    do_sample=A,    streamer=A      )

									# 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
									lowerCamelCase					:			List[str]									= cs.out[:-1]  # Remove the final "\n"
									lowerCamelCase					:			Optional[Any]									= tokenizer(A,    return_tensors='pt'      )
									self.assertEqual(streamer_text_tokenized.input_ids.shape,    (1, 1)      )




					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Union[str, Any]									= AutoTokenizer.from_pretrained('hf-internal-testing/tiny-random-gpt2'      )
									lowerCamelCase					:			List[str]									= AutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-random-gpt2'      ).to(A      )
									lowerCamelCase					:			Any									= -1

									lowerCamelCase					:			int									= ids_tensor((1, 5),    vocab_size=model.config.vocab_size      ).to(A      )
									lowerCamelCase					:			int									= TextIteratorStreamer(A,    timeout=0.001      )
									lowerCamelCase					:			Any									= {'input_ids': input_ids, 'max_new_tokens': 10, 'do_sample': False, 'streamer': streamer}
									lowerCamelCase					:			Optional[Any]									= Thread(target=model.generate,    kwargs=A      )
									thread.start()

									# The streamer will timeout after 0.001 seconds, so an exception will be raised
									with self.assertRaises(A      ):
													lowerCamelCase					:			Dict									= ''
													for new_text in streamer:
																	streamer_text += new_text


'''simple docstring'''






from manim import *




class     __snake_case		(   a__):
					def     UpperCAmelCase_							(     self      ):

									"""simple docstring"""


									lowerCamelCase					:			Optional[int]									= Rectangle(height=0.5,    width=0.5      )
									lowerCamelCase					:			List[Any]									= Rectangle(height=0.46,    width=0.46      ).set_stroke(width=0      )

									lowerCamelCase					:			List[str]									= [mem.copy() for i in range(6      )]
									lowerCamelCase					:			List[Any]									= [mem.copy() for i in range(6      )]
									lowerCamelCase					:			str									= VGroup(*A      ).arrange(A,    buff=0      )
									lowerCamelCase					:			Any									= VGroup(*A      ).arrange(A,    buff=0      )
									lowerCamelCase					:			Dict									= VGroup(A,    A      ).arrange(A,    buff=0      )
									lowerCamelCase					:			str									= Text('CPU',    font_size=24      )
									lowerCamelCase					:			int									= Group(A,    A      ).arrange(A,    buff=0.5,    aligned_edge=A      )
									cpu.move_to([-2.5, -0.5, 0]      )
									self.add(A      )

									lowerCamelCase					:			Optional[int]									= [mem.copy() for i in range(1      )]
									lowerCamelCase					:			Union[str, Any]									= VGroup(*A      ).arrange(A,    buff=0      )
									lowerCamelCase					:			Optional[Any]									= Text('GPU',    font_size=24      )
									lowerCamelCase					:			Tuple									= Group(A,    A      ).arrange(A,    buff=0.5,    aligned_edge=A      )
									gpu.align_to(A,    A      )
									gpu.set_x(gpu.get_x() - 1      )

									self.add(A      )

									lowerCamelCase					:			Optional[int]									= [mem.copy() for i in range(6      )]
									lowerCamelCase					:			Optional[Any]									= VGroup(*A      ).arrange(A,    buff=0      )

									lowerCamelCase					:			Any									= Text('Model',    font_size=24      )
									lowerCamelCase					:			Tuple									= Group(A,    A      ).arrange(A,    buff=0.5,    aligned_edge=A      )
									model.move_to([3, -1.0, 0]      )

									self.play(
									    Create(A,    run_time=1      ),    Create(A,    run_time=1      ),    Create(A,    run_time=1      ),    )

									lowerCamelCase					:			str									= MarkupText(
									    F'''First, an empty model skeleton is loaded\ninto <span fgcolor=\'{YELLOW}\'>memory</span> without using much RAM.''',    font_size=24,    )

									lowerCamelCase					:			Any									= Square(side_length=2.2      )
									key.move_to([-5, 2, 0]      )

									lowerCamelCase					:			Tuple									= MarkupText(
									    F'''<b>Key:</b>\n\n<span fgcolor=\'{YELLOW}\'>●</span> Empty Model''',    font_size=18,    )

									key_text.move_to([-5, 2.4, 0]      )

									step_a.move_to([2, 2, 0]      )
									self.play(Write(A,    run_time=2.5      ),    Write(A      ),    Write(A      )      )

									self.add(A      )

									lowerCamelCase					:			str									= []
									lowerCamelCase					:			Optional[int]									= []
									lowerCamelCase					:			Optional[Any]									= []
									for i, rect in enumerate(A      ):
													lowerCamelCase					:			List[str]									= Rectangle(height=0.46,    width=0.46      ).set_stroke(width=0.0      ).set_fill(A,    opacity=0.7      )
													cpu_target.move_to(A      )
													cpu_target.generate_target()
													lowerCamelCase					:			int									= 0.46 / 4
													lowerCamelCase					:			Optional[int]									= 0.46 / 3

													if i == 0:
																	cpu_target.target.next_to(cpu_left_col_base[0].get_corner(DOWN + LEFT      ),    buff=0.02,    direction=A      )
																	cpu_target.target.set_x(cpu_target.target.get_x() + 0.1      )
													elif i == 3:
																	cpu_target.target.next_to(cpu_targs[0].target,    direction=A,    buff=0.0      )
													else:
																	cpu_target.target.next_to(cpu_targs[i - 1].target,    direction=A,    buff=0.0      )
													cpu_targs.append(A      )

													first_animations.append(rect.animate(run_time=0.5      ).set_stroke(A      )      )
													second_animations.append(MoveToTarget(A,    run_time=1.5      )      )

									self.play(*A      )
									self.play(*A      )

									self.wait()
