





from manim import *






class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	Rectangle(height=0.5 ,						width=0.5							)
				lowerCamelCase__              =	Rectangle(height=0.46 ,						width=0.46							).set_stroke(width=0							)

				lowerCamelCase__              =	[mem.copy() for i in range(6							)]
				lowerCamelCase__              =	[mem.copy() for i in range(6							)]
				lowerCamelCase__              =	VGroup(*__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0							)
				lowerCamelCase__              =	VGroup(*__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0							)
				lowerCamelCase__              =	VGroup(__lowerCAmelCase ,						__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0							)
				lowerCamelCase__              =	Text('''CPU''' ,						font_size=2_4							)
				lowerCamelCase__              =	Group(__lowerCAmelCase ,						__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0.5 ,						aligned_edge=__lowerCAmelCase							)
				cpu.move_to([-2.5, -0.5, 0]							)
				self.add(__lowerCAmelCase							)

				lowerCamelCase__              =	[mem.copy() for i in range(4							)]
				lowerCamelCase__              =	VGroup(*__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0							)
				lowerCamelCase__              =	Text('''GPU''' ,						font_size=2_4							)
				lowerCamelCase__              =	Group(__lowerCAmelCase ,						__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0.5 ,						aligned_edge=__lowerCAmelCase							)
				gpu.move_to([-1, -1, 0]							)
				self.add(__lowerCAmelCase							)

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

				lowerCamelCase__              =	Text('''Model''' ,						font_size=2_4							)
				lowerCamelCase__              =	Group(__lowerCAmelCase ,						__lowerCAmelCase							).arrange(__lowerCAmelCase ,						buff=0.5 ,						aligned_edge=__lowerCAmelCase							)
				model.move_to([3, -1.0, 0]							)
				self.add(__lowerCAmelCase							)

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

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

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

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

				lowerCamelCase__              =	Text('''Loaded Checkpoint''' ,						font_size=2_4							)
				lowerCamelCase__              =	Group(__lowerCAmelCase ,						__lowerCAmelCase							).arrange(__lowerCAmelCase ,						aligned_edge=__lowerCAmelCase ,						buff=0.4							)
				checkpoint.move_to([3, 0.5, 0]							)

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

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

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

				self.add(__lowerCAmelCase ,						__lowerCAmelCase							)

				lowerCamelCase__              =	MarkupText(
				    F'<span fgcolor=\'{BLUE}\'>●</span> Checkpoint' ,						font_size=1_8 ,						)

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

				lowerCamelCase__              =	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=2_4 ,						)
				step_a.move_to([2, 2, 0]							)
				self.play(Write(__lowerCAmelCase							) ,						Write(__lowerCAmelCase							)							)

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

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

					lowerCamelCase__              =	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(__lowerCAmelCase ,						run_time=1.5							)							)

				self.play(*__lowerCAmelCase							)
				self.play(*__lowerCAmelCase							)
				self.wait()










from __future__ import annotations

import unittest

from transformers import EsmConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import numpy
			import tensorflow as tf

			from transformers.models.esm.modeling_tf_esm import (
			    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
			    TFEsmForMaskedLM,
			    TFEsmForSequenceClassification,
			    TFEsmForTokenClassification,
			    TFEsmModel,
			)






class       __A					:

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						):



				'''simple docstring'''







				lowerCamelCase__              =	parent
				lowerCamelCase__              =	1_3
				lowerCamelCase__              =	7
				lowerCamelCase__              =	True
				lowerCamelCase__              =	True
				lowerCamelCase__              =	True
				lowerCamelCase__              =	9_9
				lowerCamelCase__              =	3_2
				lowerCamelCase__              =	2
				lowerCamelCase__              =	4
				lowerCamelCase__              =	3_7
				lowerCamelCase__              =	'''gelu'''
				lowerCamelCase__              =	0.1
				lowerCamelCase__              =	0.1
				lowerCamelCase__              =	5_1_2
				lowerCamelCase__              =	1_6
				lowerCamelCase__              =	2
				lowerCamelCase__              =	0.02
				lowerCamelCase__              =	3
				lowerCamelCase__              =	4
				lowerCamelCase__              =	None





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size							)

				lowerCamelCase__              =	None
				if self.use_input_mask:
					lowerCamelCase__              =	random_attention_mask([self.batch_size, self.seq_length]							)

				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				if self.use_labels:
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.type_sequence_label_size							)
					lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels							)
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.num_choices							)

				lowerCamelCase__              =	EsmConfig(
				    vocab_size=self.vocab_size ,						hidden_size=self.hidden_size ,						num_hidden_layers=self.num_hidden_layers ,						pad_token_id=1 ,						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 ,						initializer_range=self.initializer_range ,						)

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





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				(
				    (
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			
				)              =	self.prepare_config_and_inputs()

				lowerCamelCase__              =	True
				lowerCamelCase__              =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
				lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						vocab_size=2							)

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModel(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	[input_ids, input_mask]
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	model(__lowerCAmelCase							)

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						):



				'''simple docstring'''







				lowerCamelCase__              =	True

				lowerCamelCase__              =	TFEsmModel(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{
				    '''input_ids''': input_ids,
				    '''attention_mask''': input_mask,
				    '''encoder_hidden_states''': encoder_hidden_states,
				    '''encoder_attention_mask''': encoder_attention_mask,
				}
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	[input_ids, input_mask]
				lowerCamelCase__              =	model(__lowerCAmelCase ,						encoder_hidden_states=__lowerCAmelCase							)

				# Also check the case where encoder outputs are not passed
				lowerCamelCase__              =	model(__lowerCAmelCase ,						attention_mask=__lowerCAmelCase							)

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmForMaskedLM(config=__lowerCAmelCase							)
				lowerCamelCase__              =	model([input_ids, input_mask]							)
				self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.num_labels
				lowerCamelCase__              =	TFEsmForTokenClassification(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
				lowerCamelCase__              =	model(__lowerCAmelCase							)
				self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels)							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.prepare_config_and_inputs()
				(
				    (
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			
				)              =	config_and_inputs
				lowerCamelCase__              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
				return config, inputs_dict






@require_tf
class       __A					(					lowerCAmelCase    ,       lowerCAmelCase    ,       unittest.TestCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			lowerCAmelCase_							  =		(
			    {
			        """feature-extraction""": TFEsmModel,
			        """fill-mask""": TFEsmForMaskedLM,
			        """text-classification""": TFEsmForSequenceClassification,
			        """token-classification""": TFEsmForTokenClassification,
			        """zero-shot""": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModelTester(self							)
				lowerCamelCase__              =	ConfigTester(self ,						config_class=__lowerCAmelCase ,						hidden_size=3_7							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				self.config_tester.run_common_tests()





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase							)





			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFEsmModel.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)





			@unittest.skip('''Protein models do not support embedding resizing.'''							)
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				pass





			@unittest.skip('''Protein models do not support embedding resizing.'''							)
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				pass





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__       ,			lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					lowerCamelCase__              =	model_class(__lowerCAmelCase							)
					assert isinstance(model.get_input_embeddings() ,						tf.keras.layers.Layer							)
					if model_class is TFEsmForMaskedLM:
						# Output embedding test differs from the main test because they're a matrix, not a layer
						lowerCamelCase__              =	model.get_bias()
						assert isinstance(__lowerCAmelCase ,						__lowerCAmelCase							)
						for k, v in name.items():
							assert isinstance(__lowerCAmelCase ,						tf.Variable							)
					else:
						lowerCamelCase__              =	model.get_output_embeddings()
						assert x is None
						lowerCamelCase__              =	model.get_bias()
						assert name is None






@require_tf
class       __A					(					unittest.TestCase				):

			'''simple docstring'''







			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

				lowerCamelCase__              =	tf.constant([[0, 1, 2, 3, 4, 5]]							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)[0]
				lowerCamelCase__              =	[1, 6, 3_3]
				self.assertEqual(list(output.numpy().shape							) ,						__lowerCAmelCase							)
				# compare the actual values for a slice.
				lowerCamelCase__              =	tf.constant(
				    [
				        [
				            [8.92_1518, -10.58_9814, -6.467_1307],
				            [-6.396_7156, -13.91_1377, -1.121_1915],
				            [-7.78_1247, -13.95_1557, -3.74_0592],
				        ]
				    ]							)
				self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,						expected_slice.numpy() ,						atol=1E-2							)							)





			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

				lowerCamelCase__              =	tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)[0]
				# compare the actual values for a slice.
				lowerCamelCase__              =	tf.constant(
				    [
				        [
				            [0.1444_3092, 0.5412_5327, 0.324_7739],
				            [0.3034_0484, 0.0052_6676, 0.3107_7722],
				            [0.3227_8043, -0.2498_7096, 0.341_4628],
				        ]
				    ]							)
				self.assertTrue(numpy.allclose(output[:, :3, :3].numpy() ,						expected_slice.numpy() ,						atol=1E-4							)							)



