Spaces:
Sleeping
Sleeping
File size: 3,171 Bytes
9218201 54d948e 9218201 54d948e 9218201 54d948e 9218201 54d948e 9218201 54d948e 9218201 54d948e 7b7a59f 54d948e 9218201 54d948e 9218201 7b7a59f 9218201 54d948e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import sys
import torch
import pytest
from pathlib import Path
from unittest.mock import patch, MagicMock
# Add repo root to sys.path
BASE_DIR = Path(__file__).resolve().parent.parent
if BASE_DIR not in sys.path:
sys.path.append(BASE_DIR)
from src.multimodal_model import MediLLMModel
BATCH_SIZE = 2
SEQ_LEN = 128
IMAGE_SIZE = (3, 224, 224)
@pytest.fixture
def dummy_inputs():
return {
"input_ids": torch.randint(0, 30522, (BATCH_SIZE, SEQ_LEN)), # dummy token IDs
"attention_mask": torch.ones(BATCH_SIZE, SEQ_LEN),
"image": torch.randn(BATCH_SIZE, *IMAGE_SIZE),
}
@patch("src.multimodal_model.AutoModel.from_pretrained")
@patch("src.multimodal_model.timm.create_model")
def test_text_only(mock_create_model, mock_auto_model, dummy_inputs):
# Mock text encoder
mock_text_encoder = MagicMock()
mock_text_encoder.config.hidden_size = 768
mock_text_encoder.return_value = MagicMock(
last_hidden_state=torch.randn(BATCH_SIZE, SEQ_LEN, 768)
)
mock_auto_model.return_value = mock_text_encoder
model = MediLLMModel(mode="text")
model.eval()
outputs = model(
input_ids=dummy_inputs["input_ids"],
attention_mask=dummy_inputs["attention_mask"]
)
assert outputs.shape == (BATCH_SIZE, 3)
probs = torch.softmax(outputs, dim=1)
assert torch.allclose(probs.sum(dim=1), torch.ones(BATCH_SIZE), atol=1e-5)
@patch("src.multimodal_model.AutoModel.from_pretrained")
@patch("src.multimodal_model.timm.create_model")
def test_image_only(mock_create_model, mock_auto_model, dummy_inputs):
# Mock image encoder
mock_image_encoder = MagicMock()
mock_image_encoder.num_features = 2048
mock_image_encoder.return_value = torch.randn(BATCH_SIZE, 2048)
mock_create_model.return_value = mock_image_encoder
model = MediLLMModel(mode="image")
model.eval()
outputs = model(image=dummy_inputs["image"])
assert outputs.shape == (BATCH_SIZE, 3)
probs = torch.softmax(outputs, dim=1)
assert torch.allclose(probs.sum(dim=1), torch.ones(BATCH_SIZE), atol=1e-5)
@patch("src.multimodal_model.AutoModel.from_pretrained")
@patch("src.multimodal_model.timm.create_model")
def test_multimodal(mock_create_model, mock_auto_model, dummy_inputs):
# Mock text encoder
mock_text_encoder = MagicMock()
mock_text_encoder.config.hidden_size = 768
mock_text_encoder.return_value = MagicMock(
last_hidden_state=torch.randn(BATCH_SIZE, SEQ_LEN, 768)
)
mock_auto_model.return_value = mock_text_encoder
# Mock image encoder
mock_image_encoder = MagicMock()
mock_image_encoder.num_features = 2048
mock_image_encoder.return_value = torch.randn(BATCH_SIZE, 2048)
mock_create_model.return_value = mock_image_encoder
model = MediLLMModel(mode="multimodal")
model.eval()
outputs = model(
input_ids=dummy_inputs["input_ids"],
attention_mask=dummy_inputs["attention_mask"],
image=dummy_inputs["image"],
)
assert outputs.shape == (BATCH_SIZE, 3)
probs = torch.softmax(outputs, dim=1)
assert torch.allclose(probs.sum(dim=1), torch.ones(BATCH_SIZE), atol=1e-5)
|