




"""simple docstring"""

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 							lowercase_			:



 '''simple docstring'''

 def __init__(					self			:			Dict       ,				_UpperCAmelCase			:			Any       ,				):
       _A     =  parent
       _A     =  13
       _A     =  7
       _A     =  True
       _A     =  True
       _A     =  True
       _A     =  99
       _A     =  32
       _A     =  2
       _A     =  4
       _A     =  37
       _A     =  'gelu'
       _A     =  0.1
       _A     =  0.1
       _A     =  512
       _A     =  16
       _A     =  2
       _A     =  0.02
       _A     =  3
       _A     =  4
       _A     =  None

 def      lowerCAmelCase_				(					self			:			Dict						):
       _A     =  ids_tensor([self.batch_size, self.seq_length]       ,				self.vocab_size						)

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

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

       _A     =  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			:			int						):
       (
           (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       
       )     =  self.prepare_config_and_inputs()

       _A     =  True
       _A     =  floats_tensor([self.batch_size, self.seq_length, self.hidden_size]						)
       _A     =  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			:			List[Any]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			int       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Optional[Any]						):
       _A     =  TFEsmModel(config=_UpperCAmelCase						)
       _A     =  {'input_ids': input_ids, 'attention_mask': input_mask}
       _A     =  model(_UpperCAmelCase						)

       _A     =  [input_ids, input_mask]
       _A     =  model(_UpperCAmelCase						)

       _A     =  model(_UpperCAmelCase						)

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

 def      lowerCAmelCase_				(					self			:			List[Any]       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Dict       ,				):
       _A     =  True

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

       _A     =  [input_ids, input_mask]
       _A     =  model(_UpperCAmelCase       ,				encoder_hidden_states=_UpperCAmelCase						)

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

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

 def      lowerCAmelCase_				(					self			:			Optional[int]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			Optional[int]       ,				_UpperCAmelCase			:			Any       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			Dict       ,				_UpperCAmelCase			:			Optional[int]						):
       _A     =  TFEsmForMaskedLM(config=_UpperCAmelCase						)
       _A     =  model([input_ids, input_mask]						)
       self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.seq_length, self.vocab_size)						)

 def      lowerCAmelCase_				(					self			:			Optional[int]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Tuple       ,				_UpperCAmelCase			:			List[str]       ,				_UpperCAmelCase			:			Tuple						):
       _A     =  self.num_labels
       _A     =  TFEsmForTokenClassification(config=_UpperCAmelCase						)
       _A     =  {'input_ids': input_ids, 'attention_mask': input_mask}
       _A     =  model(_UpperCAmelCase						)
       self.parent.assertEqual(result.logits.shape       ,				(self.batch_size, self.seq_length, self.num_labels)						)



 def      lowerCAmelCase_				(					self			:			Optional[int]						):
       _A     =  self.prepare_config_and_inputs()
       (
           (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       (
           _A
       )  ,       
       )     =  config_and_inputs
       _A     =  {'input_ids': input_ids, 'attention_mask': input_mask}
       return config, inputs_dict





@require_tf
class 							lowercase_			(	__lowerCAmelCase     ,      __lowerCAmelCase     ,      unittest.TestCase				):



 '''simple docstring'''
 UpperCAmelCase       :	List[Any]     =      (
     (
         TFEsmModel,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 UpperCAmelCase       :	str     =      (
     {
         '''feature-extraction''': TFEsmModel,
         '''fill-mask''': TFEsmForMaskedLM,
         '''text-classification''': TFEsmForSequenceClassification,
         '''token-classification''': TFEsmForTokenClassification,
         '''zero-shot''': TFEsmForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 UpperCAmelCase       :	List[str]     =      False
 UpperCAmelCase       :	str     =      False

 def      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       _A     =  TFEsmModelTester(self						)
       _A     =  ConfigTester(self       ,				config_class=_UpperCAmelCase       ,				hidden_size=37						)

 def      lowerCAmelCase_				(					self			:			Union[str, Any]						):
       self.config_tester.run_common_tests()

 def      lowerCAmelCase_				(					self			:			int						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  self.model_tester.prepare_config_and_inputs_for_decoder()
       self.model_tester.create_and_check_model_as_decoder(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_lm(*_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Dict						):
       _A     =  self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_token_classification(*_UpperCAmelCase						)

 @slow
 def      lowerCAmelCase_				(					self			:			List[str]						):
       for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
             _A     =  TFEsmModel.from_pretrained(_UpperCAmelCase						)
             self.assertIsNotNone(_UpperCAmelCase						)

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

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



 def      lowerCAmelCase_				(					self			:			Dict						):
       _A  ,       _A     =  self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
             _A     =  model_class(_UpperCAmelCase						)
             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
                   _A     =  model.get_bias()
                   assert isinstance(_UpperCAmelCase       ,				_UpperCAmelCase						)
                   for k, v in name.items():
                         assert isinstance(_UpperCAmelCase       ,				tf.Variable						)
             else:
                   _A     =  model.get_output_embeddings()
                   assert x is None
                   _A     =  model.get_bias()
                   assert name is None





@require_tf
class 							lowercase_			(	unittest.TestCase				):



 '''simple docstring'''

 @slow
 def      lowerCAmelCase_				(					self			:			List[str]						):
       _A     =  TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'						)

       _A     =  tf.constant([[0, 1, 2, 3, 4, 5]]						)
       _A     =  model(_UpperCAmelCase						)[0]
       _A     =  [1, 6, 33]
       self.assertEqual(list(output.numpy().shape						)       ,				_UpperCAmelCase						)
       # compare the actual values for a slice.
       _A     =  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			:			int						):
       _A     =  TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'						)

       _A     =  tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]						)
       _A     =  model(_UpperCAmelCase						)[0]
       # compare the actual values for a slice.
       _A     =  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						)						)







"""simple docstring"""

import os
import tempfile
import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
       import torch
       from torch import nn

       from transformers import (
           Adafactor,
           AdamW,
           get_constant_schedule,
           get_constant_schedule_with_warmup,
           get_cosine_schedule_with_warmup,
           get_cosine_with_hard_restarts_schedule_with_warmup,
           get_inverse_sqrt_schedule,
           get_linear_schedule_with_warmup,
           get_polynomial_decay_schedule_with_warmup,
       )






def 		_snake_case	(   _snake_case				:   Optional[int]							,    _snake_case				:   Optional[Any]=10				)       ->   Optional[int]:






      '''simple docstring'''


      _A     =  []
      for _ in range(_snake_case				):
            lrs.append(scheduler.get_lr()[0]				)
            scheduler.step()
      return lrs






def 		_snake_case	(   _snake_case				:   Optional[Any]							,    _snake_case				:   Union[str, Any]=10				)       ->   List[str]:






      '''simple docstring'''


      _A     =  []
      for step in range(_snake_case				):
            lrs.append(scheduler.get_lr()[0]				)
            scheduler.step()
            if step == num_steps // 2:
                  with tempfile.TemporaryDirectory() as tmpdirname:
                        _A     =  os.path.join(_snake_case							,    'schedule.bin'				)
                        torch.save(scheduler.state_dict()							,    _snake_case				)

                        _A     =  torch.load(_snake_case				)
                        scheduler.load_state_dict(_snake_case				)
      return lrs





@require_torch
class 							lowercase_			(	unittest.TestCase				):



 '''simple docstring'''

 def      lowerCAmelCase_				(					self			:			List[Any]       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			Tuple						):
       self.assertEqual(len(_UpperCAmelCase						)       ,				len(_UpperCAmelCase						)						)
       for a, b in zip(_UpperCAmelCase       ,				_UpperCAmelCase						):
             self.assertAlmostEqual(_UpperCAmelCase       ,				_UpperCAmelCase       ,				delta=_UpperCAmelCase						)

 def      lowerCAmelCase_				(					self			:			Any						):
       _A     =  torch.tensor([0.1, -0.2, -0.1]       ,				requires_grad=_UpperCAmelCase						)
       _A     =  torch.tensor([0.4, 0.2, -0.5]						)
       _A     =  nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _A     =  AdamW(params=[w]       ,				lr=2E-1       ,				weight_decay=0.0						)
       for _ in range(100						):
             _A     =  criterion(_UpperCAmelCase       ,				_UpperCAmelCase						)
             loss.backward()
             optimizer.step()
             w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
             w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()       ,				[0.4, 0.2, -0.5]       ,				tol=1E-2						)



 def      lowerCAmelCase_				(					self			:			int						):
       _A     =  torch.tensor([0.1, -0.2, -0.1]       ,				requires_grad=_UpperCAmelCase						)
       _A     =  torch.tensor([0.4, 0.2, -0.5]						)
       _A     =  nn.MSELoss()
       # No warmup, constant schedule, no gradient clipping
       _A     =  Adafactor(
           params=[w]       ,				lr=1E-2       ,				eps=(1E-3_0, 1E-3)       ,				clip_threshold=1.0       ,				decay_rate=-0.8       ,				betaa=_UpperCAmelCase       ,				weight_decay=0.0       ,				relative_step=_UpperCAmelCase       ,				scale_parameter=_UpperCAmelCase       ,				warmup_init=_UpperCAmelCase       ,				)
       for _ in range(1_000						):
             _A     =  criterion(_UpperCAmelCase       ,				_UpperCAmelCase						)
             loss.backward()
             optimizer.step()
             w.grad.detach_()  # No zero_grad() function on simple tensors. we do it ourselves.
             w.grad.zero_()
       self.assertListAlmostEqual(w.tolist()       ,				[0.4, 0.2, -0.5]       ,				tol=1E-2						)





@require_torch
class 							lowercase_			(	unittest.TestCase				):



 '''simple docstring'''
 UpperCAmelCase       :	List[str]     =      nn.Linear(50     ,      50				) if is_torch_available() else None
 UpperCAmelCase       :	Tuple     =      AdamW(m.parameters()     ,      lr=10.0				) if is_torch_available() else None
 UpperCAmelCase       :	Dict     =      10

 def      lowerCAmelCase_				(					self			:			Any       ,				_UpperCAmelCase			:			Optional[Any]       ,				_UpperCAmelCase			:			str       ,				_UpperCAmelCase			:			List[Any]       ,				_UpperCAmelCase			:			List[Any]=None						):
       self.assertEqual(len(_UpperCAmelCase						)       ,				len(_UpperCAmelCase						)						)
       for a, b in zip(_UpperCAmelCase       ,				_UpperCAmelCase						):
             self.assertAlmostEqual(_UpperCAmelCase       ,				_UpperCAmelCase       ,				delta=_UpperCAmelCase       ,				msg=_UpperCAmelCase						)



 def      lowerCAmelCase_				(					self			:			List[Any]						):
       _A     =  {'num_warmup_steps': 2, 'num_training_steps': 10}
       # schedulers doct format
       # function: (sched_args_dict, expected_learning_rates)
       _A     =  {
           get_constant_schedule: ({}, [10.0] * self.num_steps),
           get_constant_schedule_with_warmup: (
               {'num_warmup_steps': 4},
               [0.0, 2.5, 5.0, 7.5, 10.0, 10.0, 10.0, 10.0, 10.0, 10.0],
           ),
           get_linear_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 8.75, 7.5, 6.25, 5.0, 3.75, 2.5, 1.25],
           ),
           get_cosine_schedule_with_warmup: (
               {**common_kwargs},
               [0.0, 5.0, 10.0, 9.61, 8.53, 6.91, 5.0, 3.08, 1.46, 0.38],
           ),
           get_cosine_with_hard_restarts_schedule_with_warmup: (
               {**common_kwargs, 'num_cycles': 2},
               [0.0, 5.0, 10.0, 8.53, 5.0, 1.46, 10.0, 8.53, 5.0, 1.46],
           ),
           get_polynomial_decay_schedule_with_warmup: (
               {**common_kwargs, 'power': 2.0, 'lr_end': 1E-7},
               [0.0, 5.0, 10.0, 7.656, 5.625, 3.906, 2.5, 1.406, 0.625, 0.156],
           ),
           get_inverse_sqrt_schedule: (
               {'num_warmup_steps': 2},
               [0.0, 5.0, 10.0, 8.165, 7.071, 6.325, 5.774, 5.345, 5.0, 4.714],
           ),
       }

       for scheduler_func, data in scheds.items():
             _A  ,       _A     =  data

             _A     =  scheduler_func(self.optimizer       ,				**_UpperCAmelCase						)
             self.assertEqual(len([scheduler.get_lr()[0]]						)       ,				1						)
             _A     =  unwrap_schedule(_UpperCAmelCase       ,				self.num_steps						)
             self.assertListAlmostEqual(
                 _UpperCAmelCase       ,				_UpperCAmelCase       ,				tol=1E-2       ,				msg=F'''failed for {scheduler_func} in normal scheduler'''       ,				)

             _A     =  scheduler_func(self.optimizer       ,				**_UpperCAmelCase						)
             if scheduler_func.__name__ != "get_constant_schedule":
                   LambdaScheduleWrapper.wrap_scheduler(_UpperCAmelCase						)  # wrap to test picklability of the schedule
             _A     =  unwrap_and_save_reload_schedule(_UpperCAmelCase       ,				self.num_steps						)
             self.assertListEqual(_UpperCAmelCase       ,				_UpperCAmelCase       ,				msg=F'''failed for {scheduler_func} in save and reload'''						)





class 							lowercase_			:



 '''simple docstring'''

 def __init__(					self			:			Union[str, Any]       ,				_UpperCAmelCase			:			Optional[int]						):
       _A     =  fn

 def __call__(					self			:			Tuple       ,				*_UpperCAmelCase			:			List[str]       ,				**_UpperCAmelCase			:			List[str]						):
       return self.fn(*_UpperCAmelCase       ,				**_UpperCAmelCase						)



 @classmethod
 def      lowerCAmelCase_				(					self			:			Union[str, Any]       ,				_UpperCAmelCase			:			Any						):
       _A     =  list(map(self       ,				scheduler.lr_lambdas						)						)

