diff --git a/src/app/visual_bge/eva_clip/hf_model.py b/src/app/visual_bge/eva_clip/hf_model.py index c4b9fd8..51bcff3 100644 --- a/src/app/visual_bge/eva_clip/hf_model.py +++ b/src/app/visual_bge/eva_clip/hf_model.py @@ -101,8 +101,7 @@ def __init__( else: create_func, model_args = (AutoModel.from_pretrained, model_name_or_path) if pretrained else ( AutoModel.from_config, self.config) - # TODO: do all model configs have this attribute? PretrainedConfig does so yes?? - if hasattr(self.config, "is_encoder_decoder") and self.config.is_encoder_decoder: + if getattr(self.config, "is_encoder_decoder", False): self.transformer = create_func(model_args) self.transformer = self.transformer.encoder else: diff --git a/src/tests/test_hf_model.py b/src/tests/test_hf_model.py new file mode 100644 index 0000000..d2d82d7 --- /dev/null +++ b/src/tests/test_hf_model.py @@ -0,0 +1,87 @@ +import unittest +from unittest.mock import MagicMock, patch +import torch.nn as nn + +from app.visual_bge.eva_clip.hf_model import HFTextEncoder + + +class TestHFTextEncoderConfig(unittest.TestCase): + @patch("app.visual_bge.eva_clip.hf_model.AutoTokenizer") + @patch("app.visual_bge.eva_clip.hf_model.AutoModel") + @patch("app.visual_bge.eva_clip.hf_model.AutoConfig") + def test_encoder_decoder_true( + self, mock_autoconfig, mock_automodel, mock_autotokenizer + ): + mock_config = MagicMock() + mock_config.is_encoder_decoder = True + mock_config.model_type = "bert" + mock_config.hidden_size = 768 + mock_config.vocab_size = 30522 + mock_autoconfig.from_pretrained.return_value = mock_config + + mock_model_instance = MagicMock() + mock_encoder = MagicMock(spec=nn.Module) + mock_model_instance.encoder = mock_encoder + mock_automodel.from_pretrained.return_value = mock_model_instance + + encoder = HFTextEncoder( + model_name_or_path="dummy-model", + output_dim=768, + tokenizer_name="dummy-model", + ) + + mock_automodel.from_pretrained.assert_called_once_with("dummy-model") + self.assertEqual(encoder.transformer, mock_encoder) + + @patch("app.visual_bge.eva_clip.hf_model.AutoTokenizer") + @patch("app.visual_bge.eva_clip.hf_model.AutoModel") + @patch("app.visual_bge.eva_clip.hf_model.AutoConfig") + def test_encoder_decoder_false( + self, mock_autoconfig, mock_automodel, mock_autotokenizer + ): + mock_config = MagicMock() + mock_config.is_encoder_decoder = False + mock_config.model_type = "bert" + mock_config.hidden_size = 768 + mock_config.vocab_size = 30522 + mock_autoconfig.from_pretrained.return_value = mock_config + + mock_model_instance = MagicMock(spec=nn.Module) + mock_automodel.from_pretrained.return_value = mock_model_instance + + encoder = HFTextEncoder( + model_name_or_path="dummy-model", + output_dim=768, + tokenizer_name="dummy-model", + ) + + mock_automodel.from_pretrained.assert_called_once_with( + "dummy-model", add_pooling_layer=False + ) + self.assertEqual(encoder.transformer, mock_model_instance) + + @patch("app.visual_bge.eva_clip.hf_model.AutoTokenizer") + @patch("app.visual_bge.eva_clip.hf_model.AutoModel") + @patch("app.visual_bge.eva_clip.hf_model.AutoConfig") + def test_encoder_decoder_missing_attribute( + self, mock_autoconfig, mock_automodel, mock_autotokenizer + ): + mock_config = MagicMock(spec=["model_type", "hidden_size", "vocab_size"]) + mock_config.model_type = "bert" + mock_config.hidden_size = 768 + mock_config.vocab_size = 30522 + mock_autoconfig.from_pretrained.return_value = mock_config + + mock_model_instance = MagicMock(spec=nn.Module) + mock_automodel.from_pretrained.return_value = mock_model_instance + + encoder = HFTextEncoder( + model_name_or_path="dummy-model", + output_dim=768, + tokenizer_name="dummy-model", + ) + + mock_automodel.from_pretrained.assert_called_once_with( + "dummy-model", add_pooling_layer=False + ) + self.assertEqual(encoder.transformer, mock_model_instance)