Add changes for CI complicance

.gitignore

@@ -24,4 +24,13 @@ logs/
 .env 
 
 # logs
-*.log 
+*.log
+
+# --- EXCEPTIONS (ALLOW) --- 
+!data/dummy_images/
+!data/dummy_images/*.jpg
+!data/dummy_images/*.jpeg
+!data/dummy_images/*.png
+
+# Allow test CSV for CI
+!data/test_emr_records.csv 

src/generate_emr_csv.py

@@ -1,23 +1,22 @@
+import os
 import random
 import csv
 import string
 from pathlib import Path
 
+# Detect CI environment
+IS_CI = os.getenv("CI", "false").lower() == "true"
+
+# Set number of samples accordingly
+SAMPLES_PER_CLASS = 3 if IS_CI else 300  # Reduced for CI to speed up tests
+
 # Paths
 CURRENT_DIR = Path(__file__).resolve().parent
-IMAGES_DIR = CURRENT_DIR.parent / "data" / "images"
+IMAGES_DIR = CURRENT_DIR.parent / "data" / "images"  # Absolute path of images folder
 OUTPUT_FILE = CURRENT_DIR.parent / "data" / "emr_records.csv"
 
 # Label to triage
 triage_map = {"COVID": "high", "NORMAL": "low", "VIRAL PNEUMONIA": "medium"}
-SAMPLES_PER_CLASS = 300
-
-# Folders
-categories = {
-    "COVID": IMAGES_DIR / "COVID",
-    "NORMAL": IMAGES_DIR / "NORMAL",
-    "VIRAL PNEUMONIA": IMAGES_DIR / "VIRAL PNEUMONIA",
-}
 
 # Shared ambiguous templates
 shared_symptoms = [
@@ -120,7 +119,17 @@ def build_emr(label, i):
 
 
 # Generate records
-def generate_dataset():
+def generate_dataset(image_dir_override=None, output_path_override=None):
+    root_image_dir = image_dir_override or IMAGES_DIR
+    output_file = output_path_override or OUTPUT_FILE
+
+    # Folders
+    categories = {
+        "COVID": root_image_dir / "COVID",  # Absolute path of Image labels
+        "NORMAL": root_image_dir / "NORMAL",
+        "VIRAL PNEUMONIA": root_image_dir / "VIRAL PNEUMONIA",
+    }
+
     records = []
     for label, img_dir in categories.items():
         image_files = sorted(
@@ -136,7 +145,7 @@ def generate_dataset():
 
         for i in range(SAMPLES_PER_CLASS):
             image_path = str(
-                random.choice(image_files).relative_to(IMAGES_DIR.parent.parent)
+                random.choice(image_files).relative_to(root_image_dir.parent.parent)  # path of image respective to the project root
             )
             text = build_emr(label, i)
             triage = triage_map[label]
@@ -144,13 +153,13 @@ def generate_dataset():
 
     # Shuffle + write
     random.shuffle(records)
-    with open(OUTPUT_FILE, "w", newline="") as f:
+    with open(output_file, "w", newline="") as f:
         writer = csv.writer(f)
         writer.writerow(["patient_id", "image_path", "emr_text", "triage_level"])
         writer.writerows(records)
 
-    print(f"✅ Softlabel EMR dataset generated at {OUTPUT_FILE}")
+    print(f"✅ EMR dataset generated at {output_file}")
 
 
 if __name__ == "__main__":
-    generate_dataset()
+    generate_dataset(image_dir_override=IMAGES_DIR, output_path_override=OUTPUT_FILE)

src/triage_dataset.py

@@ -1,5 +1,6 @@
 import os
 import torch
+from pathlib import Path
 from torch.utils.data import Dataset
 from PIL import Image
 from torchvision import transforms
@@ -7,6 +8,9 @@ from torchvision.transforms import InterpolationMode
 import pandas as pd
 from transformers import AutoTokenizer
 
+# Check if running in CI environment
+IS_CI = os.getenv("CI", "false").lower() == "true"
+
 
 class TriageDataset(Dataset):
     def __init__(
@@ -16,6 +20,7 @@ class TriageDataset(Dataset):
         max_length=128,
         transform=None,
         mode="multimodal",
+        image_base_dir=None,
     ):
         assert mode in [
             "text",
@@ -27,6 +32,10 @@ class TriageDataset(Dataset):
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name)
         self.max_length = max_length
         self.mode = mode.lower()
+        if self.mode in ["image", "multimodal"]:
+            if image_base_dir is None:
+                raise ValueError("image directory must be provided for image or multimodal mode.")
+            self.image_base_dir = Path(image_base_dir).resolve()
 
         self.transform = (
             transform
@@ -78,11 +87,15 @@ class TriageDataset(Dataset):
 
         if self.mode in ["image", "multimodal"]:
             # Process image
-            base_dir = os.path.dirname(os.path.dirname(__file__))
-            image_path = os.path.join(base_dir, row["image_path"])
+            image_path = Path(row["image_path"])
+            if not image_path.is_absolute():
+                image_path = self.image_base_dir / image_path
 
-            if not os.path.exists(image_path):
-                raise FileNotFoundError(f"Image file not found: {image_path}")
+            if not image_path.exists():
+                if IS_CI:
+                    raise FileNotFoundError(f"[CI] Image file not found: {image_path}")
+                else:
+                    raise FileNotFoundError(f"[LOCAL] Image file not found: {image_path}")
             image = Image.open(image_path).convert("RGB")
             output["image"] = self.transform(image)
 

tests/test_generate_emr_csv.py

@@ -1,22 +1,30 @@
 import os
+import re
 import csv
 import sys
 import pytest
 from pathlib import Path
-from collections import Counter
 
-# Add repo root to the sys.path
-BASE_DIR = os.path.dirname(os.path.dirname(__file__))
-if BASE_DIR not in sys.path:
-    sys.path.append(BASE_DIR)
+# Add src/ to path so we can import from it
+BASE_DIR = Path(__file__).resolve().parent.parent
+SRC_DIR = BASE_DIR / "src"
+sys.path.insert(0, str(SRC_DIR))
 
-from src.generate_emr_csv import generate_dataset
+from generate_emr_csv import generate_dataset, OUTPUT_FILE
 
+# Determine if running in CI
+IS_CI = os.getenv("CI", "false").lower() == "true"
 
-CSV_PATH = Path(BASE_DIR) / "data" / "emr_records.csv"
-EXPECTED_CLASSES = {"low", "medium", "high"}
+# Paths
+DATA_DIR = BASE_DIR / "data"
+DUMMY_IMAGES_DIR = DATA_DIR / "dummy_images"
+REAL_IMAGES_DIR = DATA_DIR / "images"
+CSV_PATH = DATA_DIR / ("test_emr_records.csv" if IS_CI else OUTPUT_FILE)
+
+# Constants
 EXPECTED_COLUMNS = ["patient_id", "image_path", "emr_text", "triage_level"]
-EXPECTED_SAMPLES_PER_CLASS = 300
+EXPECTED_CLASSES = ["low", "medium", "high"]
+EXPECTED_SAMPLES_PER_CLASS = 3 if IS_CI else 300
 
 AMBIGUOUS_PHRASES = [
     "Symptoms could relate to a range of viral infections.",
@@ -46,54 +54,67 @@ NOISE_SENTENCES = [
 ]
 
 
-def test_dataset_generation_runs():
-    generate_dataset()
-    assert CSV_PATH.exists(), "CSV file should be generated"
-    with open(CSV_PATH, "r") as f:
-        lines = f.readlines()
-    assert len(lines) > 1  # Header + Content
+@pytest.fixture(scope="module", autouse=True)
+def generate_csv_for_test():
+    image_dir = DUMMY_IMAGES_DIR if IS_CI else REAL_IMAGES_DIR
+    generate_dataset(image_dir_override=image_dir, output_path_override=CSV_PATH)
 
 
-@pytest.fixture(scope="module")
-def load_emr_csv():
+def test_csv_exists():
     assert CSV_PATH.exists(), f"CSV file not found at: {CSV_PATH}"
+
+
+def test_csv_structure():
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.reader(f)
+        header = next(reader)
+    assert set(header) == set(EXPECTED_COLUMNS), "CSV columns mismatch"
+
+
+def test_total_and_per_class_counts():
     with open(CSV_PATH, newline="") as f:
         reader = csv.DictReader(f)
         rows = list(reader)
-    return rows
-
 
-def test_csv_structure(load_emr_csv):
-    row = load_emr_csv[0]
-    assert set(row.keys()) == set(EXPECTED_COLUMNS), "CSV columns mismatch"
+    expected_total = EXPECTED_SAMPLES_PER_CLASS * len(EXPECTED_CLASSES)
+    assert len(rows) == expected_total
 
+    counts = {"low": 0, "medium": 0, "high": 0}
+    for row in rows:
+        counts[row["triage_level"]] += 1
 
-def test_total_and_per_class_counts(load_emr_csv):
-    assert len(load_emr_csv) == 900, "Total records should be 900"
-    counts = Counter(row["triage_level"] for row in load_emr_csv)
-    for cls in EXPECTED_CLASSES:
-        assert counts[cls] == EXPECTED_SAMPLES_PER_CLASS, f"{cls} count mismatch"
+    assert all(c == EXPECTED_SAMPLES_PER_CLASS for c in counts.values)
 
 
 def test_patient_id_format_and_uniqueness(load_emr_csv):
-    ids = [row["patient_id"] for row in load_emr_csv]
-    assert all(id and "-" in id for id in ids), "Malformed patient IDs found"
-    assert len(set(ids)) == 900, "Duplicate patient IDs found"
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        ids = [row["patient_id"] for row in reader]
+        assert len(ids) == len(set(ids)), "Duplicate patient IDs found"
+        pattern = re.compile(r"^ID-[A-Z]{2}\d{2}$")
+        for pid in ids:
+            assert pattern.match(pid), f"Invalid patient ID format: {pid}"
 
 
-def test_emr_text_quality(load_emr_csv):
-    for row in load_emr_csv:
-        text = row["emr_text"]
-        assert (
-            isinstance(text, str) and len(text.split()) > 10
-        ), "EMR text too short or malformed"
-        assert "Temperature" in text and "SPO2" in text, "Vitals info missing"
+def test_emr_text_quality():
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            text = row["emr_text"]
+            assert (
+                isinstance(text, str) and len(text.split()) > 10
+            ), "EMR text too short or malformed"
+            assert "Temperature" in text and "SPO2" in text, "Vitals info missing"
 
 
-def test_image_path_format(load_emr_csv):
-    for row in load_emr_csv:
-        path = row["image_path"]
-        assert path.endswith((".jpg", ".jpeg", ".png")), f"Invalid image path: {path}"
+def test_image_path_format():
+    expected_path = DUMMY_IMAGES_DIR.relative_to(BASE_DIR) if IS_CI else REAL_IMAGES_DIR.relative_to(BASE_DIR)
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            path = row["image_path"]
+            assert path.startswith(expected_path), f"Image path should start with '{expected_path}', got: {path}"
+            assert path.endswith((".jpg", ".jpeg", ".png")), f"Invalid image path: {path}"
 
 
 def test_ambiguous_and_noise_injection(load_emr_csv):
@@ -101,28 +122,34 @@ def test_ambiguous_and_noise_injection(load_emr_csv):
     symptom_hits = 0
     noise_hits = 0
 
-    for row in load_emr_csv:
-        text = row["emr_text"]
-        if any(phrase in text for phrase in AMBIGUOUS_PHRASES):
-            ambiguous_hits += 1
-        if any(symptom in text for symptom in SHARED_SYMPTOMS):
-            symptom_hits += 1
-        if any(noise in text for noise in NOISE_SENTENCES):
-            noise_hits += 1
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            text = row["emr_text"]
+            if any(phrase in text for phrase in AMBIGUOUS_PHRASES):
+                ambiguous_hits += 1
+            if any(symptom in text for symptom in SHARED_SYMPTOMS):
+                symptom_hits += 1
+            if any(noise in text for noise in NOISE_SENTENCES):
+                noise_hits += 1
 
     assert ambiguous_hits > 800, "Ambiguous phrases missing in too many EMRs"
     assert symptom_hits > 800, "Shared symptom clues underrepresented"
     assert noise_hits > 700, "Too few EMRs contain noise sentences"
 
 
-def test_label_validity(load_emr_csv):
-    for row in load_emr_csv:
-        assert (
-            row["triage_level"] in EXPECTED_CLASSES
-        ), f"Invalid label: {row['triage_level']}"
+def test_label_validity():
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            assert (
+                row["triage_level"] in EXPECTED_CLASSES
+            ), f"Invalid label: {row['triage_level']}"
 
 
-def test_no_empty_fields(load_emr_csv):
-    for row in load_emr_csv:
-        for col in EXPECTED_COLUMNS:
-            assert row[col].strip(), f"Empty field found in colum '{col}'"
+def test_no_empty_fields():
+    with open(CSV_PATH, newline="") as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            for key, val in row.items():
+                assert val.strip() != "", f"Empty field found for {key}"

tests/test_multimodal_model.py

@@ -1,12 +1,12 @@
 import sys
-import os
 import torch
 import pytest
-from transformers import AutoTokenizer
+from pathlib import Path
+from unittest.mock import patch, MagicMock
 
 
-# Add repo root to the sys.path
-BASE_DIR = os.path.dirname(os.path.dirname(__file__))
+# 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)
 
@@ -15,57 +15,83 @@ from src.multimodal_model import MediLLMModel
 BATCH_SIZE = 2
 SEQ_LEN = 128
 IMAGE_SIZE = (3, 224, 224)
-TEXT_MODEL_NAME = "emilyalsentzer/Bio_ClinicalBERT"
-
-tokenizer = AutoTokenizer.from_pretrained(TEXT_MODEL_NAME)
 
 
 @pytest.fixture
 def dummy_inputs():
-    text_batch = ["Patient reports mild cough and fever."] * BATCH_SIZE
-    encoding = tokenizer(
-        text_batch,
-        padding="max_length",
-        truncation=True,
-        max_length=SEQ_LEN,
-        return_tensors="pt",
-    )
     return {
-        "input_ids": encoding["input_ids"],
-        "attention_mask": encoding["attention_mask"],
+        "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),
     }
 
 
-def test_text_only(dummy_inputs):
+@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"],
+        attention_mask=dummy_inputs["attention_mask"]
     )
-    assert outputs.shape == (BATCH_SIZE, 3), "Incorrect output shape for text-only mode"
 
+    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
 
-def test_image_only(dummy_inputs):
     model = MediLLMModel(mode="image")
     model.eval()
     outputs = model(image=dummy_inputs["image"])
-    assert outputs.shape == (
-        BATCH_SIZE,
-        3,
-    ), "Incorrect output shape for image-only mode"
 
+    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
 
-def test_multimodal(dummy_inputs):
     model = MediLLMModel(mode="multimodal")
     model.eval()
     outputs = model(
         input_ids=dummy_inputs["input_ids"],
-        attention_mask=dummy_inputs["attention_mask"],
+        atttention_mask=dummy_inputs["attention_mask"],
         image=dummy_inputs["image"],
     )
-    assert outputs.shape == (
-        BATCH_SIZE,
-        3,
-    ), "Incorrect output shape for multimodal mode"
+
+    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)

tests/test_triage_dataset.py

@@ -3,23 +3,36 @@ import sys
 import pytest
 import torch
 import pandas as pd
+from pathlib import Path
 
-base_dir = os.path.dirname(os.path.dirname(__file__))
-if base_dir not in sys.path:
-    sys.path.append(base_dir)
+# Setup path to src/
+BASE_DIR = Path(__file__).resolve().parent.parent
+SRC_DIR = BASE_DIR / "src"
+sys.path.insert(0, str(SRC_DIR))
 
-from src.triage_dataset import TriageDataset
+from triage_dataset import TriageDataset
 
-# Path to CSV and example image should match the local structure
-CSV_PATH = os.path.join(base_dir, "data", "emr_records.csv")
+# Detect CI environment
+IS_CI = os.getenv("CI", "false").lower() == "true"
+
+# Paths
+DATA_DIR = BASE_DIR / "data"
+CSV_PATH = DATA_DIR / ("test_emr_records.csv" if IS_CI else "emr_records.csv")
+IMAGE_DIR = DATA_DIR / ("dummy_images" if IS_CI else "images")
+EXPECTED_SAMPLES_PER_CLASS = 3 if IS_CI else 300
+EXPECTED_TOTAL = 3 * 3 if IS_CI else 300 * 3  # 3 classes
 
 
 @pytest.mark.parametrize("mode", ["text", "image", "multimodal"])
 def test_dataset_loading(mode):
-    dataset = TriageDataset(csv_file=CSV_PATH, mode=mode)
+    kwargs = {"csv_file": CSV_PATH, "mode": mode}
+    if mode in ["image", "multimodal"]:
+        kwargs["image_base_dir"] = IMAGE_DIR
+
+    dataset = TriageDataset(**kwargs)
 
     # Check dataset length
-    assert len(dataset) == 900, "Expected 900 records in the dataset"
+    assert len(dataset) == EXPECTED_TOTAL, f"Expected {EXPECTED_TOTAL} records in the dataset"
 
     # Check one sample
     sample = dataset[0]

tools/generate_dummy_images.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+import shutil
+
+# Define image categories and paths
+LABELS = ["COVID", "NORMAL", "VIRAL PNEUMONIA"]
+BASE_DIR = Path(__file__).resolve().parent.parent
+DATA_DIR = BASE_DIR / "data" / "images"
+DST_DIR = BASE_DIR / "data" / "dummy_images"
+NUM_IMAGES_PER_CLASS = 3  # keep small for CI
+
+
+def create_dummy_images():
+    for label in LABELS:
+        src_dir = DATA_DIR / label
+        dst_dir = DST_DIR / label
+        dst_dir.mkdir(parents=True, exist_ok=True)
+
+        image_files = sorted([f for f in src_dir.glob("*") if f.is_file()])
+        for i, img_path in enumerate(image_files[:NUM_IMAGES_PER_CLASS]):
+            ext = img_path.suffix
+            dummy_filename = f"dummy_{i + 1}{ext}"
+            dst_path = dst_dir / dummy_filename
+            shutil.copy(img_path, dst_path)
+
+    print(f"✅ Dummy image copies created in: {DST_DIR}")
+
+
+if __name__ == "__main__":
+    create_dummy_images()