




import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
   from transformers.models.esm.modeling_esm import (
       ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
       EsmEmbeddings,
       create_position_ids_from_input_ids,
   )

class        _a    :






     def __init__(			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE=13 ,_SCREAMING_SNAKE_CASE=7 ,_SCREAMING_SNAKE_CASE=False ,_SCREAMING_SNAKE_CASE=True ,_SCREAMING_SNAKE_CASE=False ,_SCREAMING_SNAKE_CASE=True ,_SCREAMING_SNAKE_CASE=33 ,_SCREAMING_SNAKE_CASE=32 ,_SCREAMING_SNAKE_CASE=5 ,_SCREAMING_SNAKE_CASE=4 ,_SCREAMING_SNAKE_CASE=37 ,_SCREAMING_SNAKE_CASE="gelu" ,_SCREAMING_SNAKE_CASE=0.1 ,_SCREAMING_SNAKE_CASE=0.1 ,_SCREAMING_SNAKE_CASE=512 ,_SCREAMING_SNAKE_CASE=16 ,_SCREAMING_SNAKE_CASE=2 ,_SCREAMING_SNAKE_CASE=0.0_2 ,_SCREAMING_SNAKE_CASE=3 ,_SCREAMING_SNAKE_CASE=4 ,_SCREAMING_SNAKE_CASE=None ,)							->      int:
          _snake_case							=				parent
          _snake_case							=				batch_size
          _snake_case							=				seq_length
          _snake_case							=				is_training
          _snake_case							=				use_input_mask
          _snake_case							=				use_token_type_ids
          _snake_case							=				use_labels
          _snake_case							=				vocab_size
          _snake_case							=				hidden_size
          _snake_case							=				num_hidden_layers
          _snake_case							=				num_attention_heads
          _snake_case							=				intermediate_size
          _snake_case							=				hidden_act
          _snake_case							=				hidden_dropout_prob
          _snake_case							=				attention_probs_dropout_prob
          _snake_case							=				max_position_embeddings
          _snake_case							=				type_vocab_size
          _snake_case							=				type_sequence_label_size
          _snake_case							=				initializer_range
          _snake_case							=				num_labels
          _snake_case							=				num_choices
          _snake_case							=				scope






     def     _lowercase     (			self   )							->      List[Any]:
          _snake_case							=				ids_tensor([self.batch_size, self.seq_length] ,self.vocab_size   )

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

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

          _snake_case							=				self.get_config()

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






     def     _lowercase     (			self   )							->      Any:
          return EsmConfig(
              vocab_size=self.vocab_size ,hidden_size=self.hidden_size ,pad_token_id=1 ,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 ,initializer_range=self.initializer_range ,)






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Optional[Any]:
          _snake_case							=				EsmModel(config=UpperCamelCase__   )
          model.to(UpperCamelCase__   )
          model.eval()
          _snake_case							=				model(UpperCamelCase__ ,attention_mask=UpperCamelCase__   )
          _snake_case							=				model(UpperCamelCase__   )
          _snake_case							=				model(UpperCamelCase__   )

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






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Tuple:
          _snake_case							=				EsmForMaskedLM(config=UpperCamelCase__   )
          model.to(UpperCamelCase__   )
          model.eval()
          _snake_case							=				model(UpperCamelCase__ ,attention_mask=UpperCamelCase__ ,labels=UpperCamelCase__   )
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.vocab_size)   )






     def     _lowercase     (			self ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )							->      Dict:
          _snake_case							=				self.num_labels
          _snake_case							=				EsmForTokenClassification(config=UpperCamelCase__   )
          model.to(UpperCamelCase__   )
          model.eval()
          _snake_case							=				model(UpperCamelCase__ ,attention_mask=UpperCamelCase__ ,labels=UpperCamelCase__   )
          self.parent.assertEqual(result.logits.shape ,(self.batch_size, self.seq_length, self.num_labels)   )






     def     _lowercase     (			self   )							->      Optional[Any]:
          _snake_case							=				self.prepare_config_and_inputs()
          (
              _snake_case
          )							=				config_and_inputs
          _snake_case							=				{'''input_ids''': input_ids, '''attention_mask''': input_mask}
          return config, inputs_dict







@require_torch
class        _a    (						__SCREAMING_SNAKE_CASE , __SCREAMING_SNAKE_CASE , unittest.TestCase							):
     SCREAMING_SNAKE_CASE_							:   Union[str, Any]       =						False

     SCREAMING_SNAKE_CASE_							:   int       =						(
         (
             EsmForMaskedLM,
             EsmModel,
             EsmForSequenceClassification,
             EsmForTokenClassification,
         )
         if is_torch_available()
         else ()
     )
     SCREAMING_SNAKE_CASE_							:   List[str]       =						()
     SCREAMING_SNAKE_CASE_							:   List[Any]       =						(
         {
             """feature-extraction""": EsmModel,
             """fill-mask""": EsmForMaskedLM,
             """text-classification""": EsmForSequenceClassification,
             """token-classification""": EsmForTokenClassification,
             """zero-shot""": EsmForSequenceClassification,
         }
         if is_torch_available()
         else {}
     )
     SCREAMING_SNAKE_CASE_							:   Optional[Any]       =						True






     def     _lowercase     (			self   )							->      int:
          _snake_case							=				EsmModelTester(self   )
          _snake_case							=				ConfigTester(self ,config_class=UpperCamelCase__ ,hidden_size=37   )






     def     _lowercase     (			self   )							->      Tuple:
          self.config_tester.run_common_tests()






     def     _lowercase     (			self   )							->      Tuple:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*UpperCamelCase__   )






     def     _lowercase     (			self   )							->      Any:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()
          for type in ["absolute", "relative_key", "relative_key_query"]:
               _snake_case							=				type
               self.model_tester.create_and_check_model(*UpperCamelCase__   )






     def     _lowercase     (			self   )							->      List[Any]:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__   )






     def     _lowercase     (			self   )							->      Dict:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__   )






     @slow
     def     _lowercase     (			self   )							->      Optional[int]:







          for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _snake_case							=				EsmModel.from_pretrained(UpperCamelCase__   )
               self.assertIsNotNone(UpperCamelCase__   )






     def     _lowercase     (			self   )							->      List[Any]:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()[0]
          _snake_case							=				EsmEmbeddings(config=UpperCamelCase__   )

          _snake_case							=				torch.as_tensor([[12, 31, 13, model.padding_idx]]   )
          _snake_case							=				torch.as_tensor(
              [
                  [
                      0 + model.padding_idx + 1,
                      1 + model.padding_idx + 1,
                      2 + model.padding_idx + 1,
                      model.padding_idx,
                  ]
              ]   )
          _snake_case							=				create_position_ids_from_input_ids(UpperCamelCase__ ,model.padding_idx   )
          self.assertEqual(position_ids.shape ,expected_positions.shape   )
          self.assertTrue(torch.all(torch.eq(UpperCamelCase__ ,UpperCamelCase__   )   )   )






     def     _lowercase     (			self   )							->      Optional[int]:
          _snake_case							=				self.model_tester.prepare_config_and_inputs()[0]
          _snake_case							=				EsmEmbeddings(config=UpperCamelCase__   )

          _snake_case							=				torch.empty(2 ,4 ,30   )
          _snake_case							=				[
              0 + embeddings.padding_idx + 1,
              1 + embeddings.padding_idx + 1,
              2 + embeddings.padding_idx + 1,
              3 + embeddings.padding_idx + 1,
          ]
          _snake_case							=				torch.as_tensor([expected_single_positions, expected_single_positions]   )
          _snake_case							=				embeddings.create_position_ids_from_inputs_embeds(UpperCamelCase__   )
          self.assertEqual(position_ids.shape ,expected_positions.shape   )
          self.assertTrue(torch.all(torch.eq(UpperCamelCase__ ,UpperCamelCase__   )   )   )






     @unittest.skip("Esm does not support embedding resizing"   )
     def     _lowercase     (			self   )							->      int:
          pass






     @unittest.skip("Esm does not support embedding resizing"   )
     def     _lowercase     (			self   )							->      str:
          pass






     @unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."   )
     def     _lowercase     (			self   )							->      Optional[int]:
          pass







@require_torch
class        _a    (						__SCREAMING_SNAKE_CASE							):






     @slow
     def     _lowercase     (			self   )							->      Union[str, Any]:







          with torch.no_grad():
               _snake_case							=				EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"   )
               model.eval()
               _snake_case							=				torch.tensor([[0, 1, 2, 3, 4, 5]]   )
               _snake_case							=				model(UpperCamelCase__   )[0]

               _snake_case							=				33

               _snake_case							=				torch.Size((1, 6, vocab_size)   )
               self.assertEqual(output.shape ,UpperCamelCase__   )

               _snake_case							=				torch.tensor(
                   [[[8.9_2_1_5, -1_0.5_8_9_8, -6.4_6_7_1], [-6.3_9_6_7, -1_3.9_1_1_4, -1.1_2_1_2], [-7.7_8_1_2, -1_3.9_5_1_6, -3.7_4_0_6]]]   )
               self.assertTrue(torch.allclose(output[:, :3, :3] ,UpperCamelCase__ ,atol=1e-4   )   )






     @slow
     def     _lowercase     (			self   )							->      str:







          with torch.no_grad():
               _snake_case							=				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"   )
               model.eval()

               _snake_case							=				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]   )
               _snake_case							=				model(UpperCamelCase__   )[0]
               # compare the actual values for a slice.
               _snake_case							=				torch.tensor(
                   [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]   )
               self.assertTrue(torch.allclose(output[:, :3, :3] ,UpperCamelCase__ ,atol=1e-4   )   )








'''simple docstring'''





from manim import *

class        _a    (						__lowerCAmelCase							):






     def     _lowercase     (			self   )							->      Optional[int]:
          _snake_case							=				Rectangle(height=0.5 ,width=0.5   )
          _snake_case							=				Rectangle(height=0.4_6 ,width=0.4_6   ).set_stroke(width=0   )

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

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

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

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

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

               _snake_case							=				Rectangle(height=0.4_6 / 4 ,width=0.4_6 / 3   ).set_stroke(width=0.0   ).set_fill(_SCREAMING_SNAKE_CASE ,opacity=0.7   )

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

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

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

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

          _snake_case							=				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(_SCREAMING_SNAKE_CASE ,_SCREAMING_SNAKE_CASE   )

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

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

          _snake_case							=				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(_SCREAMING_SNAKE_CASE   ) ,Write(_SCREAMING_SNAKE_CASE   )   )

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

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

               _snake_case							=				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(_SCREAMING_SNAKE_CASE ,run_time=1.5   )   )

          self.play(*_SCREAMING_SNAKE_CASE   )
          self.play(*_SCREAMING_SNAKE_CASE   )
          self.wait()


