



from __future__ import annotations

import unittest

from transformers import is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow


if is_tf_available():
						import tensorflow as tf

						from transformers import AutoTokenizer, TFAutoModelForSeqaSeqLM






@require_tf
@require_sentencepiece
@require_tokenizers
class      UpperCamelCase    (						unittest.TestCase     ):



	@slow
	def 			__A							(       self		):
							A__							=       TFAutoModelForSeqaSeqLM.from_pretrained("google/mt5-small"		)
							A__							=       AutoTokenizer.from_pretrained("google/mt5-small"		)

							A__							=       tokenizer("Hello there"			,	return_tensors="tf"		).input_ids
							A__							=       tokenizer("Hi I am"			,	return_tensors="tf"		).input_ids

							A__							=       model(UpperCAmelCase__			,	labels=UpperCAmelCase__		).loss
							A__							=       -tf.math.reduce_mean(UpperCAmelCase__		).numpy()

							A__							=       -21.228_168
							self.assertTrue(abs(mtf_score - EXPECTED_SCORE		) < 2e-4		)










from manim import *






class      UpperCamelCase    (						_UpperCAmelCase     ):



	def 			__A							(       self		):
							A__							=       Rectangle(height=0.5			,	width=0.5		)
							A__							=       Rectangle(height=0.46			,	width=0.46		).set_stroke(width=0		)

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       VGroup(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       Text("CPU"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							cpu.move_to([-2.5, -0.5, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       [mem.copy() for i in range(4		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)
							A__							=       Text("GPU"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							gpu.move_to([-1, -1, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)

							A__							=       Text("Model"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0.5			,	aligned_edge=UpperCAmelCase__		)
							model.move_to([3, -1.0, 0]		)
							self.add(UpperCAmelCase__		)

							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													rect.set_stroke(UpperCAmelCase__		)
													# target = fill.copy().set_fill(YELLOW, opacity=0.7)
													# target.move_to(rect)
													# self.add(target)

													A__							=       Rectangle(height=0.46 / 4			,	width=0.46 / 3		).set_stroke(width=0.0		).set_fill(UpperCAmelCase__			,	opacity=0.7		)

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

							A__							=       [mem.copy() for i in range(6		)]
							A__							=       VGroup(*UpperCAmelCase__		).arrange(UpperCAmelCase__			,	buff=0		)

							A__							=       Text("Loaded Checkpoint"			,	font_size=24		)
							A__							=       Group(UpperCAmelCase__			,	UpperCAmelCase__		).arrange(UpperCAmelCase__			,	aligned_edge=UpperCAmelCase__			,	buff=0.4		)
							checkpoint.move_to([3, 0.5, 0]		)

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

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

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

							self.add(UpperCAmelCase__			,	UpperCAmelCase__		)

							A__							=       MarkupText(
							    F"""<span fgcolor='{BLUE}'>●</span> Checkpoint"""			,	font_size=18			,	)

							blue_text.next_to(UpperCAmelCase__			,	DOWN * 2.4			,	aligned_edge=key_text.get_left()		)

							A__							=       MarkupText(
							    F"""Next, a <i><span fgcolor=\"{BLUE}\">second</span></i> model is loaded into memory,\nwith the weights of a <span fgcolor=\"{BLUE}\">single shard</span>."""			,	font_size=24			,	)
							step_a.move_to([2, 2, 0]		)
							self.play(Write(UpperCAmelCase__		)			,	Write(UpperCAmelCase__		)		)

							self.play(Write(UpperCAmelCase__			,	run_time=1		)			,	Create(UpperCAmelCase__			,	run_time=1		)		)

							A__							=       []
							A__							=       []
							for i, rect in enumerate(UpperCAmelCase__		):
													A__							=       fill.copy().set_fill(UpperCAmelCase__			,	opacity=0.7		)
													target.move_to(UpperCAmelCase__		)
													first_animations.append(GrowFromCenter(UpperCAmelCase__			,	run_time=1		)		)

													A__							=       target.copy()
													cpu_target.generate_target()
													if i < 5:
																			cpu_target.target.move_to(cpu_left_col_base[i + 1]		)
													else:
																			cpu_target.target.move_to(cpu_right_col_base[i - 5]		)
													second_animations.append(MoveToTarget(UpperCAmelCase__			,	run_time=1.5		)		)

							self.play(*UpperCAmelCase__		)
							self.play(*UpperCAmelCase__		)
							self.wait()





