Add tests, formatting, and CI enhancements with pre-commit support

.flake8

@@ -0,0 +1,11 @@
+# .flake8
+[flake8]
+max-line-length = 88
+ignore = E501, E402
+exclude = 
+    .git,
+    __pycache__,
+    .venv,
+    env,
+    build,
+    dist

.github/workflows/ci.yml

@@ -23,12 +23,19 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -r requirements.txt
-          pip install pytest flake8
+          pip install pytest flake8 black isort
 
-      - name: ✅ Lint code
+      - name: ✅ Lint code with flake8
+        run: flake8
+
+      - name: 🔧 Check code format with black
+        run: |
+          black --check .
+
+      - name: 📦 Check import order with isort
         run: |
-          flake8 src/ --ignore=E501
+          isort . --check-only
 
       - name: 🧪 Run unit tests
         run: |
-          pytest tests/
+          pytest --cov=src tests/

.gitignore

@@ -3,6 +3,7 @@ data/
 checkpoints/
 __pycache__/
 *.py[cod]
+.coverage
 
 # Weights & Biases
 wandb/

.pre-commit-config.yaml

@@ -0,0 +1,7 @@
+repos:
+  - repo: https://github.com/pycqa/flake8
+    rev: 6.1.0
+    hooks:
+      - id: flake8
+        additional_dependencies: []
+        args: ["--ignore=E501,E402"]

experiments/train_optuna.py

@@ -14,7 +14,7 @@ from torch.utils.data import DataLoader, Subset
 from torch.nn import CrossEntropyLoss
 from torch.optim import Adam 
 from sklearn.model_selection import StratifiedShuffleSplit
-from sklearn.metrics import accuracy_score, f1_score, classification_report, confusion_matrix
+from sklearn.metrics import accuracy_score, f1_score, confusion_matrix
 
 
 # Setup base path

pyproject.toml

@@ -0,0 +1,6 @@
+[tool.black]
+line-length = 88
+
+[tool.isort]
+profile = "black"
+line_length = 88

requirements.txt

@@ -36,4 +36,6 @@ python-multipart>0.0.6
 
 # Linting and testing
 pytest>=7.4.0
+pytest-cov>=4.1
+pre-commit>=3.5.0
 flake8>=6.1.0

src/generate_emr_csv.py

@@ -29,7 +29,7 @@ shared_symptoms = [
     "Vital signs mostly stable; slight variation in temperature.",
 ]
 
-# Overlapping diagnosis clues
+# Overlapping diagnosis clues to add ambiguity
 shared_diagnosis = [
     "Symptoms could relate to a range of viral infections.",
     "Presentation not distinctly matching any single infection.",
@@ -120,28 +120,39 @@ def build_emr(label, i):
 
 
 # Generate records
-records = []
-for label, img_dir in categories.items():
-    image_files = sorted(
-        [
-            f
-            for f in img_dir.glob("*")
-            if f.suffix.lower() in [".png", ".jpg", ".jpeg"]
-        ]
-    )
-    for i in range(SAMPLES_PER_CLASS):
-        image_path = str(
-            random.choice(image_files).relative_to(IMAGES_DIR.parent.parent)
+def generate_dataset():
+    records = []
+    for label, img_dir in categories.items():
+        image_files = sorted(
+            [
+                f
+                for f in img_dir.glob("*")
+                if f.suffix.lower() in [".png", ".jpg", ".jpeg"]
+            ]
         )
-        text = build_emr(label, i)
-        triage = triage_map[label]
-        records.append([f"{label}-{i+1}", image_path, text, triage])
-
-# Shuffle + write
-random.shuffle(records)
-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}")
+        for i in range(SAMPLES_PER_CLASS):
+            image_path = str(
+                random.choice(image_files)
+                .relative_to(IMAGES_DIR.parent.parent)
+            )
+            text = build_emr(label, i)
+            triage = triage_map[label]
+            records.append([f"{label}-{i+1}", image_path, text, triage])
+
+    # Shuffle + write
+    random.shuffle(records)
+    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}")
+
+
+if __name__ == "__main__":
+    generate_dataset()

tests/test_generate_emr_csv.py

@@ -0,0 +1,131 @@
+import os
+import csv
+import sys
+import pytest
+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)
+
+from src.generate_emr_csv import generate_dataset, OUTPUT_FILE
+
+
+CSV_PATH = OUTPUT_FILE
+EXPECTED_CLASSES = {"low", "medium", "high"}
+EXPECTED_COLUMNS = ["patient_id", "image_path", "emr_text", "triage_level"]
+EXPECTED_SAMPLES_PER_CLASS = 300
+
+AMBIGUOUS_PHRASES = [
+    "Symptoms could relate to a range of viral infections.",
+    "Presentation not distinctly matching any single infection.",
+    "Further tests required to confirm diagnosis.",
+    "Findings are borderline; clinical judgment advised.",
+    "Observation warranted due to overlapping signs.",
+    "Initial assessment inconclusive.",
+]
+
+SHARED_SYMPTOMS = [
+    "Mild cough and slight fever reported.",
+    "General fatigue and throat irritation present.",
+    "Breathing mildly labored during physical exertion.",
+    "No major respiratory distress; mild wheezing noted.",
+    "Occasional chest tightness reported.",
+    "Vital signs mostly stable; slight variation in temperature.",
+]
+
+NOISE_SENTENCES = [
+    "Patient is cooperative and alert.",
+    "Dietary habits unremarkable.",
+    "Hydration status normal.",
+    "Follow-up advised if symptoms persist.",
+    "No notable family medical history.",
+    "No medications currently administered.",
+]
+
+
+def test_dataset_generation_runs():
+    generate_dataset()
+    assert CSV_PATH.exists(), "CSV file should be generated"
+    with open(OUTPUT_FILE, "r") as f:
+        lines = f.readlines()
+    assert len(lines) > 1  # Header + Content
+
+
+@pytest.fixture(scope="module")
+def load_emr_csv():
+    assert CSV_PATH.exists(), f"CSV file not found at: {CSV_PATH}"
+    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"
+
+
+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"
+        )
+
+
+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"
+
+
+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_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_ambiguous_and_noise_injection(load_emr_csv):
+    ambiguous_hits = 0
+    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
+
+    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_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}'"

tests/test_multimodal_model.py

@@ -0,0 +1,73 @@
+import sys
+import os
+import torch
+import pytest
+from transformers import AutoTokenizer
+
+
+# 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)
+
+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"],
+        "image": torch.randn(BATCH_SIZE, *IMAGE_SIZE),
+    }
+
+
+def test_text_only(dummy_inputs):
+    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), (
+        "Incorrect output shape for text-only mode"
+    )
+
+
+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"
+
+
+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"],
+        image=dummy_inputs["image"],
+    )
+    assert outputs.shape == (
+        BATCH_SIZE,
+        3,
+    ), "Incorrect output shape for multimodal mode"

tests/test_triage_dataset.py

@@ -0,0 +1,41 @@
+import os
+import sys
+import pytest
+import torch
+
+base_dir = os.path.dirname(os.path.dirname(__file__))
+if base_dir not in sys.path:
+    sys.path.append(base_dir)
+
+from src.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")
+
+
+@pytest.mark.parametrize("mode", ["text", "image", "multimodal"])
+def test_dataset_loading(mode):
+    dataset = TriageDataset(csv_file=CSV_PATH, mode=mode)
+
+    # Check dataset length
+    assert len(dataset) == 900, "Expected 900 records in the dataset"
+
+    # Check one sample
+    sample = dataset[0]
+
+    if mode in ["text", "multimodal"]:
+        assert "input_ids" in sample, (
+            "Missing input_ids in text/multimodal mode"
+        )
+        assert (
+            "attention_mask" in sample
+        ), "Missing attention_mask in text/multimodal mode"
+        assert sample["input_ids"].shape[0] == 128, "Incorrect token length"
+
+    if mode in ["image", "multimodal"]:
+        assert "image" in sample, "Missing image in image/multimodal mode"
+        assert isinstance(sample["image"], torch.Tensor), "Image not a tensor"
+        assert sample["image"].shape[1:] == (224, 224), "Incorrect image size"
+
+    assert "label" in sample, "Missing label"
+    assert sample["label"].item() in [0, 1, 2], "Invalid label value"