Medi-LLM: public demo

app/demo/demo.py

@@ -70,7 +70,7 @@ def classify(role, mode, normalize_mode, emr_text, image, use_rollout):
 
     # Model caching
     if mode not in model_cache:
-        model_cache[mode] = load_model(mode, MODEL_PATHS[mode])
+        model_cache[mode] = load_model(mode)
     model = model_cache[mode]
 
     # Run prediction
@@ -529,301 +529,306 @@ def reset_ui():
     )
 
 
-# --- Gradio UI ---
-style_path = Path(__file__).resolve().parent / "style.css"
-with open(style_path, "r") as f:
-    custom_css = f.read()
-
-with gr.Blocks(css=custom_css) as demo:
-    # Centered title and subtitle
-    gr.Markdown("<h2 class='centered'>🩺 Medi-LLM: Clinical Triage Assistant 🩻</h2>")
-    gr.Markdown("<p class='centered'>Upload a chest X-ray and/or enter EMR text to get a triage level prediction.</p>")
-    gr.HTML(
-        """
-        <div class='welcome-banner' style="background-color: #24283b; border-left: 4px solid #7aa2f7; padding: 16px; border-radius: 8px; margin-bottom: 16px;">
-        <h3 style="margin-top: 0; color: #c0caf5;">👋 Welcome to Medi-LLM</h3>
-        <p style="color: #a9b1d6; line-height: 1.6;">
-            This AI assistant helps triage patients using <strong>EMR text</strong> and <strong>chest X-rays</strong>.<br>
-            📝 Enter EMR notes, 📷 upload a chest X-ray, or use both for a multimodal diagnosis.<br>
-            👩‍⚕️ Select <strong>Doctor</strong> mode to view insights like Grad-CAM heatmaps and token-level attention.<br>
-            💾 Save your results for later by exporting them to a CSV file.
-        </p>
-        </div>
-        """
-    )
-
-    # Hidden State
-    role_state = gr.State(value="User")
-    mode_state = gr.State(value=DEFAULT_MODE)
-    rollout_state = gr.State(value=False)
-    normaliza_mode_state = gr.State(value="visual")
-    inference_done = gr.State(value=False)
-
-    # Role and Mode selection
-    with gr.Row(equal_height=True):
-        with gr.Column():
-            role = gr.Radio(["User", "Doctor"], value="User", label="Select Role", info="Doctors see insights like Grad-CAM and token attention", elem_id="role_selector")
-            mode = gr.Radio(["text", "image", "multimodal"], value=DEFAULT_MODE, label="Select Input Mode", info="Choose Diagnosis input type", elem_id="mode_selector")
-            with gr.Column(visible=False) as normalize_mode_column:
-                normalize_mode = gr.Radio(
-                    ["visual", "probabilistic"],
-                    value="visual",
-                    label="Attention Normalization",
-                    info="Softmax sums to 1 (probabilistic). Visual uses gamma-boosted scaling for color clarity."
-                )
-                use_rollout = gr.Checkbox(
-                    label="Use attention rollout (CLS -> inputs)",
-                    value=False,
-                    info="Includes residuals and multiplies attention across layers. Slower but often more faithful."
-                )
-
-    normalize_mode.change(
-        fn=lambda val: val,
-        inputs=[normalize_mode],
-        outputs=[normaliza_mode_state]
-    )
-
-    use_rollout.change(
-        fn=lambda v: v,
-        inputs=[use_rollout],
-        outputs=[rollout_state]
-    )
+def build_ui():
+    # Load CSS safely (don't crash if file is missing on remote)
+    style_path = Path(__file__).resolve().parent / "style.css"
+    custom_css = style_path.read_text(encoding="utf-8") if style_path.exists() else ""
+
+    with gr.Blocks(css=custom_css) as demo:
+        # ----- Header -----
+        gr.Markdown("<h2 class='centered'>🩺 Medi-LLM: Clinical Triage Assistant 🩻</h2>")
+        gr.Markdown("<p class='centered'>Upload a chest X-ray and/or enter EMR text to get a triage level prediction.</p>")
+        gr.HTML(
+            """
+            <div class='welcome-banner' style="background-color: #24283b; border-left: 4px solid #7aa2f7; padding: 16px; border-radius: 8px; margin-bottom: 16px;">
+             <h3 style="margin-top: 0; color: #c0caf5;">👋 Welcome to Medi-LLM</h3>
+             <p style="color: #a9b1d6; line-height: 1.6;">
+               This AI assistant helps triage patients using <strong>EMR text</strong> and <strong>chest X-rays</strong>.<br>
+               📝 Enter EMR notes, 📷 upload a chest X-ray, or use both for a multimodal diagnosis.<br>
+               👩‍⚕️ Select <strong>Doctor</strong> mode to view insights like Grad-CAM heatmaps and token-level attention.<br>
+               💾 Save your results for later by exporting them to a CSV file.
+             </p>
+            </div>
+            """
+        )
 
-    # Input: EMR text and/or image
-    with gr.Row():
-        with gr.Column(scale=3, elem_id="text_col") as text_col:
-            emr_text, image, max_file_note = render_inputs(DEFAULT_MODE)
+        # ----- Hidden State -----
+        role_state = gr.State(value="User")
+        mode_state = gr.State(value=DEFAULT_MODE)
+        rollout_state = gr.State(value=False)
+        normaliza_mode_state = gr.State(value="visual")
+        inference_done = gr.State(value=False)
+
+        # ----- Role and Mode selection -----
+        with gr.Row(equal_height=True):
+            with gr.Column():
+                role = gr.Radio(["User", "Doctor"], value="User", label="Select Role", info="Doctors see insights like Grad-CAM and token attention", elem_id="role_selector")
+                mode = gr.Radio(["text", "image", "multimodal"], value=DEFAULT_MODE, label="Select Input Mode", info="Choose Diagnosis input type", elem_id="mode_selector")
+                with gr.Column(visible=False) as normalize_mode_column:
+                    normalize_mode = gr.Radio(
+                        ["visual", "probabilistic"],
+                        value="visual",
+                        label="Attention Normalization",
+                        info="Softmax sums to 1 (probabilistic). Visual uses gamma-boosted scaling for color clarity."
+                    )
+                    use_rollout = gr.Checkbox(
+                        label="Use attention rollout (CLS -> inputs)",
+                        value=False,
+                        info="Includes residuals and multiplies attention across layers. Slower but often more faithful."
+                    )
+
+        # ----- Inputs -----
+        with gr.Row():
+            with gr.Column(scale=3, elem_id="text_col"):
+                emr_text, image, max_file_note = render_inputs(DEFAULT_MODE)
+
+        # ----- Actions -----
+        with gr.Row():
+            submit_btn = gr.Button(
+                "🔍 Run Inference",
+                elem_id="inference_btn"
+            )
+            reset_btn = gr.Button(
+                "↩️ Reset",
+                elem_id="reset_btn"
+            )
 
-    # Submit button
-    with gr.Row():
-        submit_btn = gr.Button(
-            "🔍 Run Inference",
-            elem_id="inference_btn"
+        # ----- Outputs -----
+        with gr.Column(elem_classes=["output-box"]):
+            result_box = gr.Textbox(label="🧪 Triage Prediction", interactive=False)
+            confidence_label = gr.Label(label="📊 Confidence", visible=False)
+            prediction_count_box = gr.Textbox(value="Predictions: 0", interactive=False, label="🧮 Count", elem_id="prediction_count_box")
+            insights_tab = gr.Tabs(visible=False)
+            class_probs_json = gr.JSON(label="🔍 Class Probabilities", visible=True, elem_classes=["json-box"])
+            with insights_tab:
+                with gr.Tab("📷 Grad-CAM"):
+                    gradcam_img = gr.Image(visible=False, elem_classes=["gr-image-box"])
+                with gr.Tab("🔬 Token Attention"):
+                    token_attention = gr.HighlightedText(
+                        visible=False,
+                        show_legend=False,
+                        color_map={
+                            "0.0": "#7aa2f7",   # blue
+                            "0.25": "#80deea",  # cyan
+                            "0.5": "#fbc02d",   # yellow
+                            "0.75": "#ff8a65",  # orange
+                            "1.0": "#f7768e",   # red
+                        },
+                        elem_classes=["token-attn-box"]
+                    )
+                    top5_html = gr.HTML(value="", visible=False)
+
+                    inject_tooltips()
+
+                    gr.HTML("""
+                    <div class="attention-legend">
+                        <div style="display: flex; align-items: center; gap: 8px;">
+                            <span style="font-size: 14px; color: #c0caf5;">0.0</span>
+                            <div class="attention-gradient-bar"></div>
+                            <span style="font-size: 14px; color: #c0caf5;">1.0</span>
+                        </div>
+                    </div>
+                    """)
+
+        with gr.Row():
+            loading_msg = gr.Markdown(value="", visible=False, elem_classes=["loading-msg"])
+
+        # ----- Inference Wiring -----
+        submit_btn.click(
+            fn=show_loading_msg,
+            outputs=[loading_msg]
+        ).then(
+            fn=classify,
+            inputs=[role_state, mode_state, normaliza_mode_state, emr_text, image, rollout_state],
+            outputs=[
+                result_box,
+                gradcam_img,
+                token_attention,
+                top5_html,
+                confidence_label,
+                insights_tab,
+                prediction_count_box,
+                class_probs_json,
+            ]
+        ).then(
+            fn=lambda: gr.update(value="", visible=False),
+            outputs=[loading_msg]
+        ).then(
+            fn=lambda: True,
+            outputs=[inference_done]
         )
-        reset_btn = gr.Button(
-            "↩️ Reset",
-            elem_id="reset_btn"
+
+        # ----- Role/Mode/Param Change Wiring -----
+        role.change(
+            fn=update_role_state,
+            inputs=[role],
+            outputs=[role_state, normalize_mode_column, insights_tab, token_attention, gradcam_img, use_rollout, top5_html]
         )
 
-    # Outputs
-    with gr.Column(elem_classes=["output-box"]):
-        result_box = gr.Textbox(label="🧪 Triage Prediction", interactive=False)
-        confidence_label = gr.Label(label="📊 Confidence", visible=False)
-        prediction_count_box = gr.Textbox(value="Predictions: 0", interactive=False, label="🧮 Count", elem_id="prediction_count_box")
-        insights_tab = gr.Tabs(visible=False)
-        class_probs_json = gr.JSON(label="🔍 Class Probabilities", visible=True, elem_classes=["json-box"])
-        with insights_tab:
-            with gr.Tab("📷 Grad-CAM"):
-                gradcam_img = gr.Image(visible=False, elem_classes=["gr-image-box"])
-            with gr.Tab("🔬 Token Attention"):
-                token_attention = gr.HighlightedText(
-                    visible=False,
-                    show_legend=False,
-                    color_map={
-                        "0.0": "#7aa2f7",   # blue
-                        "0.25": "#80deea",  # cyan
-                        "0.5": "#fbc02d",   # yellow
-                        "0.75": "#ff8a65",  # orange
-                        "1.0": "#f7768e",   # red
-                    },
-                    elem_classes=["token-attn-box"]
-                )
-                top5_html = gr.HTML(value="", visible=False)
+        # Input Updates
+        mode.change(
+            fn=lambda m: (*render_inputs(m), m),
+            inputs=[mode],
+            outputs=[emr_text, image, max_file_note, mode_state]
+        )
 
-                inject_tooltips()
+        normalize_mode.change(
+            fn=lambda val: val,
+            inputs=[normalize_mode],
+            outputs=[normaliza_mode_state]
+        )
 
-                gr.HTML("""
-                <div class="attention-legend">
-                    <div style="display: flex; align-items: center; gap: 8px;">
-                        <span style="font-size: 14px; color: #c0caf5;">0.0</span>
-                        <div class="attention-gradient-bar"></div>
-                        <span style="font-size: 14px; color: #c0caf5;">1.0</span>
-                    </div>
-                </div>
-                """)
-
-    with gr.Row():
-        loading_msg = gr.Markdown(value="", visible=False, elem_classes=["loading-msg"])
-
-    # Bind inference
-    submit_btn.click(
-        fn=show_loading_msg,
-        outputs=[loading_msg]
-    ).then(
-        fn=classify,
-        inputs=[role_state, mode_state, normaliza_mode_state, emr_text, image, rollout_state],
-        outputs=[
-            result_box,
-            gradcam_img,
-            token_attention,
-            top5_html,
-            confidence_label,
-            insights_tab,
-            prediction_count_box,
-            class_probs_json,
-        ]
-    ).then(
-        fn=lambda: gr.update(value="", visible=False),
-        outputs=[loading_msg]
-    ).then(
-        fn=lambda: True,
-        outputs=[inference_done]
-    )
+        use_rollout.change(
+            fn=lambda v: v,
+            inputs=[use_rollout],
+            outputs=[rollout_state]
+        )
 
-    # Input Updates
-    mode.change(
-        fn=lambda m: (*render_inputs(m), m),
-        inputs=[mode],
-        outputs=[emr_text, image, max_file_note, mode_state]
-    )
+        normalize_mode.change(
+            fn=rerun_if_done,
+            inputs=[inference_done, role_state, mode_state, normalize_mode, emr_text, image, rollout_state],
+            outputs=[
+                result_box,
+                gradcam_img,
+                token_attention,
+                top5_html,
+                confidence_label,
+                insights_tab,
+                prediction_count_box,
+                class_probs_json,
+            ]
+        )
 
-    role.change(
-        fn=update_role_state,
-        inputs=[role],
-        outputs=[role_state, normalize_mode_column, insights_tab, token_attention, gradcam_img, use_rollout, top5_html]
-    )
+        use_rollout.change(
+            fn=rerun_if_done,
+            inputs=[inference_done, role_state, mode_state, normalize_mode, emr_text, image, rollout_state],
+            outputs=[
+                result_box,
+                gradcam_img,
+                token_attention,
+                top5_html,
+                confidence_label,
+                insights_tab,
+                prediction_count_box,
+                class_probs_json
+            ]
+        )
 
-    normalize_mode.change(
-        fn=rerun_if_done,
-        inputs=[inference_done, role_state, mode_state, normalize_mode, emr_text, image, rollout_state],
-        outputs=[
-            result_box,
-            gradcam_img,
-            token_attention,
-            top5_html,
-            confidence_label,
-            insights_tab,
-            prediction_count_box,
-            class_probs_json,
-        ]
-    )
+        # ----- CSV Export & Log Controls -----
+        gr.Markdown("### 📁 Export Prediction Log")
 
-    use_rollout.change(
-        fn=rerun_if_done,
-        inputs=[inference_done, role_state, mode_state, normalize_mode, emr_text, image, rollout_state],
-        outputs=[
-            result_box,
-            gradcam_img,
-            token_attention,
-            top5_html,
-            confidence_label,
-            insights_tab,
-            prediction_count_box,
-            class_probs_json
-        ]
-    )
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=3):
+                filename_input = gr.Textbox(
+                    label="CSV filename (optional)",
+                    placeholder="e.g., triage_results.csv",
+                    info="Set filename as needed or leave blank for auto-naming",
+                    elem_id="csv_filename"
+                )
 
-    # CSV Export UI
-    gr.Markdown("### 📁 Export Prediction Log")
+                export_status_box = gr.Textbox(
+                    value="",
+                    visible=False,
+                    interactive=False,
+                    label="",
+                    elem_id="export_status"
+                )
 
-    with gr.Row(equal_height=True):
-        with gr.Column(scale=3):
-            filename_input = gr.Textbox(
-                label="CSV filename (optional)",
-                placeholder="e.g., triage_results.csv",
-                info="Set filename as needed or leave blank for auto-naming",
-                elem_id="csv_filename"
-            )
+            with gr.Column(scale=4):
+                gr.Markdown(
+                    "📑 **Summary**\n\nDownload your triage results for clinical review or research.",
+                    elem_classes="centered"
+                )
+                with gr.Row():
+                    with gr.Column(scale=1, min_width=200):
+                        download_btn = gr.Button("💾 Export CSV", elem_id="export_button")
+                    with gr.Column(scale=1, min_width=200):
+                        clear_btn = gr.Button("🗑️ Clear Logs", elem_id="clear_button")
+                confirm_clear_btn = gr.Button("✅ Confirm Clear", visible=False, elem_id="confirm_button")
+                confirm_box = gr.Textbox(label="Status", interactive=False, visible=False, elem_id="confirm_box")
+
+            with gr.Column(scale=3):
+                csv_output = gr.File(label="📂 Download Link", elem_id="download_box")
+
+        download_btn.click(
+            fn=export_csv,
+            inputs=[filename_input, role_state],
+            outputs=[
+                csv_output,
+                csv_output,
+                export_status_box
+            ]
+        ).then(
+            fn=blink_box_effect,
+            inputs=[csv_output],
+            outputs=[csv_output]
+        ).then(
+            fn=disable_filename_input,
+            outputs=[filename_input]
+        )
 
-            export_status_box = gr.Textbox(
-                value="",
-                visible=False,
-                interactive=False,
-                label="",
-                elem_id="export_status"
-            )
+        clear_btn.click(
+            fn=lambda: (
+                confirm_clear(),
+                gr.Button(visible=True),
+            ),
+            outputs=[confirm_box, confirm_clear_btn]
+        )
 
-        with gr.Column(scale=4):
-            gr.Markdown(
-                "📑 **Summary**\n\nDownload your triage results for clinical review or research.",
-                elem_classes="centered"
-            )
-            with gr.Row():
-                with gr.Column(scale=1, min_width=200):
-                    download_btn = gr.Button("💾 Export CSV", elem_id="export_button")
-                with gr.Column(scale=1, min_width=200):
-                    clear_btn = gr.Button("🗑️ Clear Logs", elem_id="clear_button")
-            confirm_clear_btn = gr.Button("✅ Confirm Clear", visible=False, elem_id="confirm_button")
-            confirm_box = gr.Textbox(label="Status", interactive=False, visible=False, elem_id="confirm_box")
-
-        with gr.Column(scale=3):
-            csv_output = gr.File(label="📂 Download Link", elem_id="download_box")
-
-    download_btn.click(
-        fn=export_csv,
-        inputs=[filename_input, role_state],
-        outputs=[
-            csv_output,
-            csv_output,
-            export_status_box
-        ]
-    ).then(
-        fn=blink_box_effect,
-        inputs=[csv_output],
-        outputs=[csv_output]
-    ).then(
-        fn=disable_filename_input,
-        outputs=[filename_input]
-    )
+        confirm_clear_btn.click(
+            fn=clear_confirmed,
+            inputs=[role_state],
+            outputs=[
+                prediction_count_box,  # reset prediction count
+                confirm_box,           # show success message
+                csv_output,            # hide CSV output file
+                filename_input         # re-enable input box
+            ]
+        ).then(
+            fn=lambda: gr.update(visible=False),  # Hide confirm button
+            outputs=[confirm_clear_btn]
+        ).then(
+            fn=reset_confirm_box,
+            outputs=[confirm_box]
+        )
 
-    clear_btn.click(
-        fn=lambda: (
-            confirm_clear(),
-            gr.Button(visible=True),
-        ),
-        outputs=[confirm_box, confirm_clear_btn]
-    )
+        # ----- Reset Wiring -----
+        reset_btn.click(
+            fn=reset_ui,
+            outputs=[
+                emr_text,               # 1
+                image,                  # 2
+                max_file_note,          # 3
+                result_box,             # 4
+                gradcam_img,            # 5
+                token_attention,        # 6
+                top5_html,              # 7
+                confidence_label,       # 8
+                insights_tab,           # 9
+                class_probs_json,       # 10
+                role_state,             # 11
+                mode_state,             # 12
+                normaliza_mode_state,   # 13
+                role,                   # 14 (radio)
+                mode,                   # 15 (radio)
+                normalize_mode,         # 16 (radio)
+                normalize_mode_column,  # 17 (column visibility)
+                use_rollout,            # 18
+                rollout_state,          # 19
+                loading_msg,            # 20
+                inference_done,         # 21
+                export_status_box       # 22
+            ]
+        )
+    return demo
 
-    confirm_clear_btn.click(
-        fn=clear_confirmed,
-        inputs=[role_state],
-        outputs=[
-            prediction_count_box,  # reset prediction count
-            confirm_box,           # show success message
-            csv_output,            # hide CSV output file
-            filename_input         # re-enable input box
-        ]
-    ).then(
-        fn=lambda: gr.update(visible=False),  # Hide confirm button
-        outputs=[confirm_clear_btn]
-    ).then(
-        fn=reset_confirm_box,
-        outputs=[confirm_box]
-    )
 
-    # Reset UI
-    reset_btn.click(
-        fn=reset_ui,
-        outputs=[
-            emr_text,               # 1
-            image,                  # 2
-            max_file_note,          # 3
-            result_box,             # 4
-            gradcam_img,            # 5
-            token_attention,        # 6
-            top5_html,              # 7
-            confidence_label,       # 8
-            insights_tab,           # 9
-            class_probs_json,       # 10
-            role_state,             # 11
-            mode_state,             # 12
-            normaliza_mode_state,   # 13
-            role,                   # 14 (radio)
-            mode,                   # 15 (radio)
-            normalize_mode,         # 16 (radio)
-            normalize_mode_column,  # 17 (column visibility)
-            use_rollout,            # 18
-            rollout_state,          # 19
-            loading_msg,            # 20
-            inference_done,         # 21
-            export_status_box       # 22
-        ]
-    )
+# Expose for Spaces & imports
+demo = build_ui()
 
 if __name__ == "__main__":
-    for mode, path in MODEL_PATHS.items():
-        if not os.path.exists(path):
-            print(f"❌ Missing model for mode {mode}: {path}")
-            print("Please download or train your models before launching the demo.")
-            exit(1)
-    demo.launch()
+    demo.launch(
+        server_name=os.getenv("GRADIO_SERVER_NAME", "127.0.0.1"),
+        server_port=int(os.getenv("GRADIO_SERVER_PORT", "7860")),
+        show_error=True,
+    )

app/utils/create_hf_space.py

@@ -0,0 +1,9 @@
+from huggingface_hub import HfApi
+HfApi().create_repo(
+    repo_id="Preetham22/medi-llm",
+    repo_type="space",
+    space_sdk="gradio",
+    exist_ok=True,
+)
+
+print("✅ Space ready: Preetham22/medi-llm")

app/utils/inference_utils.py

@@ -1,3 +1,4 @@
+import os
 import sys
 import torch
 import yaml
@@ -5,6 +6,7 @@ import numpy as np
 from pathlib import Path
 from transformers import AutoTokenizer
 from torchvision import transforms
+from huggingface_hub import hf_hub_download
 
 ROOT_DIR = Path(__file__).resolve().parent.parent.parent
 sys.path.append(str(ROOT_DIR))
@@ -12,10 +14,42 @@ sys.path.append(str(ROOT_DIR))
 from src.multimodal_model import MediLLMModel
 from app.utils.gradcam_utils import register_hooks, generate_gradcam
 
-
+# --------------------
+# Runtime / Hub config
+# --------------------
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# Label map
+# Map modes -> filenames in  HF model repo
+HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "Preetham22/medi-llm-weights")
+HF_WEIGHTS_REV = os.getenv("HF_WEIGHTS_REV", None)  # optional (commit/tag/branch)
+FILENAMES = {
+    "text": "medi_llm_state_dict_text.pth",
+    "image": "medi_llm_state_dict_image.pth",
+    "multimodal": "medi_llm_state_dict_multimodal.pth"
+}
+
+
+def resolve_weights_path(mode: str) -> str:
+    """Download (or reuse cached)  weights for the given mode from HF Hub."""
+    if mode not in FILENAMES:
+        raise ValueError(f"Unknown mode '{mode}'. Expected one of {list(FILENAMES)}.")
+    filename = FILENAMES[mode]
+    try:
+        return hf_hub_download(
+            repo_id=HF_MODEL_REPO,
+            filename=filename,
+            revision=HF_WEIGHTS_REV  # can be None
+        )
+    except Exception as e:
+        raise RuntimeError(
+            f"Failed to fetch weights '{filename}' from repo '{HF_MODEL_REPO}'."
+            f"Set HF_MODEL_REPO or check filenames. Original error: {e}"
+        )
+
+
+# ----------------------
+# Labels / preprocessing
+# ----------------------
 inv_map = {0: "low", 1: "medium", 2: "high"}
 
 # Tokenizer and image transform
@@ -27,22 +61,63 @@ image_transform = transforms.Compose([
 ])
 
 
-def load_model(mode, model_path, config_path=str(Path("config/config.yaml").resolve())):
+# ----------------------
+# Model load
+# ----------------------
+def _safe_torch_load(path: str, map_location: torch.device):
+    """
+    Prefer weights_only=True (newer Pytorch), but fall back if not supported.
+    """
+    try:
+        return torch.load(path, map_location=map_location, weights_only=True)  # PyTorch >= 2.2/2.3
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+
+
+def load_model(mode: str, config_path: str = str(Path("config/config.yaml").resolve())):
+    """
+    Load MediLLMModel for the given mode and populate weights from HF Hub.
+    Expects config/config.yaml with keys per mode (dropout, hidden_dim).
+    """
     with open(config_path, "r") as f:
-        config = yaml.safe_load(f)[mode]
+        cfg_all = yaml.safe_load(f)
+    if mode not in cfg_all:
+        raise KeyError(f"Mode '{mode}' not found in {config_path}. Keys: {list(cfg_all.keys())}")
+    config = cfg_all[mode]
 
+    # Build model
     model = MediLLMModel(
         mode=mode,
         dropout=config["dropout"],
         hidden_dim=config["hidden_dim"]
     )
-    state = torch.load(model_path, map_location=DEVICE)
-    model.load_state_dict(state)
+
+    # Download weights & load
+    weights_path = resolve_weights_path(mode)
+    state = _safe_torch_load(weights_path, map_location=DEVICE)
+
+    # Sometimes checkpoints save as {'state_dict': ...}
+    if isinstance(state, dict) and "state_dict" in state:
+        state = state["state_dict"]
+
+    try:
+        model.load_state_dict(state)  # strict by default
+    except RuntimeError as e:
+        # allow non-strict if minor mismatches (buffer names)
+        try:
+            model.load_state_dict(state, strict=False)
+            print(f"⚠️ Loaded with strict=False due to: {e}")
+        except Exception:
+            raise
+
     model.to(DEVICE)
     model.eval()
     return model
 
 
+# -----------------------
+# Attention rollout utils
+# -----------------------
 def attention_rollout(attentions, last_k=4, residual_alpha=0.5):
     """
     attentions_tuple: tuple/list of layer attentions; each is (B,H,S,S)
@@ -139,6 +214,9 @@ def _normalize_for_display_wordlevel(attn_scores, normalize_mode="visual", tempe
         return attn_array0, labels
 
 
+# ------------------
+# Prediction
+# ------------------
 def predict(
     model,
     mode,

requirements-dev.txt

@@ -0,0 +1,54 @@
+# (Bring in everything you listed, with safer caps)
+
+# PyTorch stack (match minor versions!)
+torch>=2.2,<2.4
+torchvision>=0.17,<0.19
+torchaudio>=2.2,<2.4
+
+# Transformers & NLP
+transformers>=4.41,<4.46
+datasets>=2.19,<2.21
+accelerate>=0.28,<0.31
+peft>=0.11,<0.13
+tokenizers>=0.15     # usually pulled by transformers, but explicit helps
+safetensors>=0.4.3
+huggingface_hub>=0.24,<0.26
+# If you ever use T5/ALBERT etc., add:
+# sentencepiece>=0.1.99
+
+# Vision models
+timm>=0.9.7,<1.0
+opencv-python-headless>=4.8
+Pillow>=10.0,<11
+
+# Optimization / tracking
+optuna>=3.5,<4
+wandb>=0.16.6,<0.18
+
+# Eval / viz
+scikit-learn>=1.3,<1.6
+matplotlib>=3.8,<3.9
+seaborn>=0.13,<0.14
+tqdm>=4.66,<5
+
+# Data processing
+pandas>=2.2,<2.3
+numpy>=1.26,<2.2
+pyyaml>=6.0
+scipy>=1.11,<1.14   # sklearn relies on it; make it explicit to avoid surprises
+
+# Optional deployment (FastAPI)
+fastapi>=0.110,<0.114
+pydantic>=2.5,<3
+uvicorn>=0.27,<0.31
+python-multipart>=0.0.6
+# Optional perf:
+json>=3.9
+
+# Linting & testing
+pytest>=7.4,<9
+pytest-cov>=4.1,<5
+pre-commit>=3.5,<4
+flake8>=6.1,<7
+# Optional modern linter:
+ruff>=0.4,<0.7

requirements.txt

@@ -1,41 +1,19 @@
-# Core PyTorch stack (CPU or GPU version to be installed separately)
-torch>=2.1.0
-torchvision>=0.16.0
-torchaudio>=2.1.0
+# Core runtime (CPU)
+torch>=2.2,<2.4
+torchvision>=0.17,<0.19
+# torchaudio not needed for this app; add if you really use it
 
-# Transformers and NLP
-transformers>=4.35.0
-datasets>=2.14.0
-accelerate>=0.25.0
-peft>=0.9.0
+# UI + model fetch
+gradio>=3.45.2,<3.47
+huggingface_hub>=0.24,<0.26
+safetensors>=0.4.3
 
-# Vision and image models
-opencv-python>=4.8.0
-Pillow>=10.0.0
-timm>=0.9.2
-
-# Optimization and hyperparameter tuning
-optuna>=3.3.0
-wandb>=0.15.0
-
-# Evaluation and visualization
-scikit-learn>=1.3.0
-matplotlib>=3.8.0
-seaborn>=0.13.0
-tqdm>=4.65.0
-
-# Data processing
-pandas>=2.1.0 
-numpy>=1.25.0
+# Image / utils
+opencv-python-headless>=4.8
+Pillow>=10.0,<11
+pandas>=2.2,<2.3
+numpy>=1.26,<2.2
+scikit-learn>=1.3,<1.6
+tqdm>=4.66,<5
+matplotlib>=3.8,<3.9
 pyyaml>=6.0
-
-# FastAPI for deployment
-fastapi>=0.100.0
-uvicorn>=0.27.0
-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