





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



 def __init__(				self,__lowerCamelCase,):
        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 							UpperCamelCase     (				self					):
        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 							UpperCamelCase     (				self					):
        (
            (
            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 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
        A__		   =		TFEsmModel(config=__lowerCamelCase					)
        A__		   =		{'''input_ids''': input_ids, '''attention_mask''': input_mask}
        A__		   =		model(__lowerCamelCase					)

        A__		   =		[input_ids, input_mask]
        A__		   =		model(__lowerCamelCase					)

        A__		   =		model(__lowerCamelCase					)

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



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,):
        A__		   =		True

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

        A__		   =		[input_ids, input_mask]
        A__		   =		model(__lowerCamelCase,encoder_hidden_states=__lowerCamelCase					)

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

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



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
        A__		   =		TFEsmForMaskedLM(config=__lowerCamelCase					)
        A__		   =		model([input_ids, input_mask]					)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.seq_length, self.vocab_size)					)



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




 def 							UpperCamelCase     (				self					):
        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        SCREAMING_SNAKE_CASE__					(		UpperCamelCase__							,   UpperCamelCase__							,   unittest.TestCase ):
 __SCREAMING_SNAKE_CASE							=      (
     (
         TFEsmModel,
         TFEsmForMaskedLM,
         TFEsmForSequenceClassification,
         TFEsmForTokenClassification,
     )
     if is_tf_available()
     else ()
 )
 __SCREAMING_SNAKE_CASE							=      (
     {
         '''feature-extraction''': TFEsmModel,
         '''fill-mask''': TFEsmForMaskedLM,
         '''text-classification''': TFEsmForSequenceClassification,
         '''token-classification''': TFEsmForTokenClassification,
         '''zero-shot''': TFEsmForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __SCREAMING_SNAKE_CASE							=      False
 __SCREAMING_SNAKE_CASE							=      False



 def 							UpperCamelCase     (				self					):
        A__		   =		TFEsmModelTester(self					)
        A__		   =		ConfigTester(self,config_class=__lowerCamelCase,hidden_size=37					)



 def 							UpperCamelCase     (				self					):
        self.config_tester.run_common_tests()



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs_for_decoder()
        self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_masked_lm(*__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_token_classification(*__lowerCamelCase					)



 @slow
 def 							UpperCamelCase     (				self					):
        for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               A__		   =		TFEsmModel.from_pretrained(__lowerCamelCase					)
               self.assertIsNotNone(__lowerCamelCase					)



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



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




 def 							UpperCamelCase     (				self					):
        A__   ,      A__		   =		self.model_tester.prepare_config_and_inputs_for_common()

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




@require_tf
class        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 @slow
 def 							UpperCamelCase     (				self					):
        A__		   =		TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''					)

        A__		   =		tf.constant([[0, 1, 2, 3, 4, 5]]					)
        A__		   =		model(__lowerCamelCase					)[0]
        A__		   =		[1, 6, 33]
        self.assertEqual(list(output.numpy().shape					),__lowerCamelCase					)
        # compare the actual values for a slice.
        A__		   =		tf.constant(
            [
                [
                    [8.921518, -10.589814, -6.4671307],
                    [-6.3967156, -13.911377, -1.1211915],
                    [-7.781247, -13.951557, -3.740592],
                ]
            ]					)
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy(),expected_slice.numpy(),atol=1E-2					)					)




 @slow
 def 							UpperCamelCase     (				self					):
        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(__lowerCamelCase					)[0]
        # compare the actual values for a slice.
        A__		   =		tf.constant(
            [
                [
                    [0.14443092, 0.54125327, 0.3247739],
                    [0.30340484, 0.00526676, 0.31077722],
                    [0.32278043, -0.24987096, 0.3414628],
                ]
            ]					)
        self.assertTrue(numpy.allclose(output[:, :3, :3].numpy(),expected_slice.numpy(),atol=1E-4					)					)












import shutil
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_tf_cross_test,
    require_tf,
    require_torch,
    require_torchvision,
    require_vision,
)
from transformers.utils import is_tf_available, is_torch_available, is_vision_available


if is_vision_available():
       from PIL import Image

       from transformers import AutoProcessor, SamImageProcessor, SamProcessor

if is_torch_available():
       import torch

if is_tf_available():
       import tensorflow as tf




@require_vision
@require_torchvision
class        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 def 							UpperCamelCase     (				self					):
        A__		   =		tempfile.mkdtemp()
        A__		   =		SamImageProcessor()
        A__		   =		SamProcessor(__lowerCamelCase					)
        processor.save_pretrained(self.tmpdirname					)



 def 							UpperCamelCase     (				self,**__lowerCamelCase					):
        return AutoProcessor.from_pretrained(self.tmpdirname,**__lowerCamelCase					).image_processor



 def 							UpperCamelCase     (				self					):
        shutil.rmtree(self.tmpdirname					)



 def 							UpperCamelCase     (				self					):
        A__		   =		[np.random.randint(255,size=(3, 30, 400),dtype=np.uinta					)]

        A__		   =		[Image.fromarray(np.moveaxis(__lowerCamelCase,0,-1					)					) for x in image_inputs]

        return image_inputs



 def 							UpperCamelCase     (				self					):
        A__		   =		SamProcessor(image_processor=self.get_image_processor()					)
        processor.save_pretrained(self.tmpdirname					)

        A__		   =		self.get_image_processor(do_normalize=__lowerCamelCase,padding_value=1.0					)

        A__		   =		SamProcessor.from_pretrained(self.tmpdirname,do_normalize=__lowerCamelCase,padding_value=1.0					)

        self.assertEqual(processor.image_processor.to_json_string(),image_processor_add_kwargs.to_json_string()					)
        self.assertIsInstance(processor.image_processor,__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)

        A__		   =		self.prepare_image_inputs()

        A__		   =		image_processor(__lowerCamelCase,return_tensors='''np'''					)
        A__		   =		processor(images=__lowerCamelCase,return_tensors='''np'''					)

        input_feat_extract.pop('''original_sizes'''					)  # pop original_sizes as it is popped in the processor
        input_feat_extract.pop('''reshaped_input_sizes'''					)  # pop original_sizes as it is popped in the processor

        for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum(),input_processor[key].sum(),delta=1E-2					)




 @require_torch
 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)
        A__		   =		[torch.ones((1, 3, 5, 5)					)]

        A__		   =		[[1764, 2646]]

        A__		   =		[[683, 1024]]
        A__		   =		processor.post_process_masks(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					)
        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        A__		   =		processor.post_process_masks(
            __lowerCamelCase,torch.tensor(__lowerCamelCase					),torch.tensor(__lowerCamelCase					)					)
        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        # should also work with np
        A__		   =		[np.ones((1, 3, 5, 5)					)]
        A__		   =		processor.post_process_masks(__lowerCamelCase,np.array(__lowerCamelCase					),np.array(__lowerCamelCase					)					)

        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        A__		   =		[[1, 0], [0, 1]]
        with self.assertRaises(__lowerCamelCase					):
               A__		   =		processor.post_process_masks(__lowerCamelCase,np.array(__lowerCamelCase					),np.array(__lowerCamelCase					)					)




@require_vision
@require_tf
class        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 def 							UpperCamelCase     (				self					):
        A__		   =		tempfile.mkdtemp()
        A__		   =		SamImageProcessor()
        A__		   =		SamProcessor(__lowerCamelCase					)
        processor.save_pretrained(self.tmpdirname					)



 def 							UpperCamelCase     (				self,**__lowerCamelCase					):
        return AutoProcessor.from_pretrained(self.tmpdirname,**__lowerCamelCase					).image_processor



 def 							UpperCamelCase     (				self					):
        shutil.rmtree(self.tmpdirname					)



 def 							UpperCamelCase     (				self					):
        A__		   =		[np.random.randint(255,size=(3, 30, 400),dtype=np.uinta					)]

        A__		   =		[Image.fromarray(np.moveaxis(__lowerCamelCase,0,-1					)					) for x in image_inputs]

        return image_inputs



 def 							UpperCamelCase     (				self					):
        A__		   =		SamProcessor(image_processor=self.get_image_processor()					)
        processor.save_pretrained(self.tmpdirname					)

        A__		   =		self.get_image_processor(do_normalize=__lowerCamelCase,padding_value=1.0					)

        A__		   =		SamProcessor.from_pretrained(self.tmpdirname,do_normalize=__lowerCamelCase,padding_value=1.0					)

        self.assertEqual(processor.image_processor.to_json_string(),image_processor_add_kwargs.to_json_string()					)
        self.assertIsInstance(processor.image_processor,__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)

        A__		   =		self.prepare_image_inputs()

        A__		   =		image_processor(__lowerCamelCase,return_tensors='''np'''					)
        A__		   =		processor(images=__lowerCamelCase,return_tensors='''np'''					)

        input_feat_extract.pop('''original_sizes'''					)  # pop original_sizes as it is popped in the processor
        input_feat_extract.pop('''reshaped_input_sizes'''					)  # pop reshaped_input_sizes as it is popped in the processor

        for key in input_feat_extract.keys():
               self.assertAlmostEqual(input_feat_extract[key].sum(),input_processor[key].sum(),delta=1E-2					)




 @require_tf
 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)
        A__		   =		[tf.ones((1, 3, 5, 5)					)]

        A__		   =		[[1764, 2646]]

        A__		   =		[[683, 1024]]
        A__		   =		processor.post_process_masks(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase,return_tensors='''tf'''					)
        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        A__		   =		processor.post_process_masks(
            __lowerCamelCase,tf.convert_to_tensor(__lowerCamelCase					),tf.convert_to_tensor(__lowerCamelCase					),return_tensors='''tf''',)
        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        # should also work with np
        A__		   =		[np.ones((1, 3, 5, 5)					)]
        A__		   =		processor.post_process_masks(
            __lowerCamelCase,np.array(__lowerCamelCase					),np.array(__lowerCamelCase					),return_tensors='''tf'''					)

        self.assertEqual(masks[0].shape,(1, 3, 1764, 2646)					)

        A__		   =		[[1, 0], [0, 1]]
        with self.assertRaises(tf.errors.InvalidArgumentError					):
               A__		   =		processor.post_process_masks(
                   __lowerCamelCase,np.array(__lowerCamelCase					),np.array(__lowerCamelCase					),return_tensors='''tf'''					)




@require_vision
@require_torchvision
class        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 def 							UpperCamelCase     (				self					):
        A__		   =		tempfile.mkdtemp()
        A__		   =		SamImageProcessor()
        A__		   =		SamProcessor(__lowerCamelCase					)
        processor.save_pretrained(self.tmpdirname					)



 def 							UpperCamelCase     (				self,**__lowerCamelCase					):
        return AutoProcessor.from_pretrained(self.tmpdirname,**__lowerCamelCase					).image_processor



 def 							UpperCamelCase     (				self					):
        shutil.rmtree(self.tmpdirname					)



 def 							UpperCamelCase     (				self					):
        A__		   =		[np.random.randint(255,size=(3, 30, 400),dtype=np.uinta					)]

        A__		   =		[Image.fromarray(np.moveaxis(__lowerCamelCase,0,-1					)					) for x in image_inputs]

        return image_inputs



 @is_pt_tf_cross_test
 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)
        A__		   =		np.random.randint(0,2,size=(1, 3, 5, 5)					).astype(np.floataa					)
        A__		   =		[tf.convert_to_tensor(__lowerCamelCase					)]
        A__		   =		[torch.tensor(__lowerCamelCase					)]

        A__		   =		[[1764, 2646]]

        A__		   =		[[683, 1024]]
        A__		   =		processor.post_process_masks(
            __lowerCamelCase,__lowerCamelCase,__lowerCamelCase,return_tensors='''tf'''					)
        A__		   =		processor.post_process_masks(
            __lowerCamelCase,__lowerCamelCase,__lowerCamelCase,return_tensors='''pt'''					)

        self.assertTrue(np.all(tf_masks[0].numpy() == pt_masks[0].numpy()					)					)




 @is_pt_tf_cross_test
 def 							UpperCamelCase     (				self					):
        A__		   =		self.get_image_processor()

        A__		   =		SamProcessor(image_processor=__lowerCamelCase					)

        A__		   =		self.prepare_image_inputs()

        A__		   =		image_processor(__lowerCamelCase,return_tensors='''pt'''					)['''pixel_values'''].numpy()
        A__		   =		processor(images=__lowerCamelCase,return_tensors='''pt'''					)['''pixel_values'''].numpy()

        A__		   =		image_processor(__lowerCamelCase,return_tensors='''tf'''					)['''pixel_values'''].numpy()
        A__		   =		processor(images=__lowerCamelCase,return_tensors='''tf'''					)['''pixel_values'''].numpy()

        self.assertTrue(np.allclose(__lowerCamelCase,__lowerCamelCase					)					)
        self.assertTrue(np.allclose(__lowerCamelCase,__lowerCamelCase					)					)
        self.assertTrue(np.allclose(__lowerCamelCase,__lowerCamelCase					)					)





