add loggs

notebooks/kaggle_medic_demo.ipynb

@@ -12,16 +12,16 @@
     "|---|---|---|\n",
     "| 1 · Intake Historian | Patient data, CrCl, MDR risk | MedGemma 4B IT |\n",
     "| 2 · Vision Specialist | Lab report → structured JSON | MedGemma 4B IT |\n",
-    "| 3 · Trend Analyst | MIC creep, resistance velocity | MedGemma 27B Text IT ¹ |\n",
-    "| 4 · Clinical Pharmacologist | Final Rx + safety check | MedGemma 4B IT + TxGemma 9B ¹ |\n",
-    "\n",
-    "> ¹ Substituted with smaller variants on Kaggle T4 (16 GB GPU) — see Section 3.\n",
+    "| 3 · Trend Analyst | MIC creep, resistance velocity | MedGemma 27B Text IT |\n",
+    "| 4 · Clinical Pharmacologist | Final Rx + safety check | MedGemma 4B IT + TxGemma 9B |\n",
     "\n",
     "**Before running this notebook:**\n",
     "1. Click **Add data** (top-right) → search for **`mghobashy/drug-drug-interactions`** → add it\n",
     "2. Add your HuggingFace token under **Add-ons → Secrets** as `HF_TOKEN`\n",
     "3. Accept model licences on HuggingFace (see Section 2)\n",
     "\n",
+    "**Requirements:** GPU with ≥ 30 GB VRAM (e.g. Kaggle P100 / A100 / T4 ×2)\n",
+    "\n",
     "**Steps:** Clone → Install → Authenticate → Download models → Init KB → Launch app"
    ]
   },
@@ -69,7 +69,12 @@
    "id": "4c637bc0",
    "metadata": {},
    "outputs": [],
-   "source": "%%bash\n# Always start fresh to avoid stale code from previous runs\nrm -rf /kaggle/working/AMR-Guard\ngit clone \"$GITHUB_REPO\" /kaggle/working/AMR-Guard"
+   "source": [
+    "%%bash\n",
+    "# Always start fresh to avoid stale code from previous runs\n",
+    "rm -rf /kaggle/working/AMR-Guard\n",
+    "git clone \"$GITHUB_REPO\" /kaggle/working/AMR-Guard"
+   ]
   },
   {
    "cell_type": "code",
@@ -98,7 +103,8 @@
     "\n",
     "Accept the model licences **before** running this notebook:\n",
     "- MedGemma 4B IT → https://huggingface.co/google/medgemma-4b-it\n",
-    "- TxGemma 2B → https://huggingface.co/google/txgemma-2b-predict"
+    "- MedGemma 27B Text IT → https://huggingface.co/google/medgemma-27b-text-it\n",
+    "- TxGemma 9B → https://huggingface.co/google/txgemma-9b-predict"
    ]
   },
   {
@@ -130,15 +136,13 @@
    "source": [
     "## 3 · Download Models\n",
     "\n",
-    "| Model | Agent | VRAM (4-bit) | Kaggle T4 |\n",
-    "|---|---|---|---|\n",
-    "| `google/medgemma-4b-it` | 1, 2, 4 primary | ~3 GB | ✓ |\n",
-    "| `google/medgemma-27b-text-it` | 3 (Trend Analyst) | ~14 GB | marginal — using 4B sub |\n",
-    "| `google/txgemma-9b-predict` | 4 safety check | ~5 GB | ✓ (stacked with 4B: ~8 GB) |\n",
-    "| `google/txgemma-2b-predict` | 4 safety fallback | ~1.5 GB | ✓ |\n",
+    "| Model | Agent | VRAM (4-bit) |\n",
+    "|---|---|---|\n",
+    "| `google/medgemma-4b-it` | 1, 2, 4 primary | ~3 GB |\n",
+    "| `google/medgemma-27b-text-it` | 3 (Trend Analyst) | ~14 GB |\n",
+    "| `google/txgemma-9b-predict` | 4 safety check | ~5 GB |\n",
     "\n",
-    "**Kaggle strategy:** download `medgemma-4b-it` and `txgemma-2b-predict`.  \n",
-    "Swap to the full 27B / 9B on a machine with ≥ 24 GB VRAM by editing the variables below."
+    "**Total estimated VRAM:** ~22 GB in 4-bit quantization."
    ]
   },
   {
@@ -224,7 +228,33 @@
    "id": "a61f1fb1",
    "metadata": {},
    "outputs": [],
-   "source": "# Write .env\nenv = f\"\"\"\nMEDIC_ENV=kaggle\nMEDIC_QUANTIZATION=4bit\n\n# Agent 1, 2, 4 — MedGemma 4B IT\nMEDIC_LOCAL_MEDGEMMA_4B_MODEL={MEDGEMMA_4B}\n\n# Agent 3 — MedGemma 27B Text IT  (subbed with 4B for Kaggle T4)\n# To use full 27B: set to google/medgemma-27b-text-it\nMEDIC_LOCAL_MEDGEMMA_27B_MODEL={MEDGEMMA_4B}\n\n# Agent 4 safety — TxGemma 9B  (subbed with 2B for Kaggle T4)\n# To use full 9B: set to google/txgemma-9b-predict\nMEDIC_LOCAL_TXGEMMA_9B_MODEL={TXGEMMA_2B}\nMEDIC_LOCAL_TXGEMMA_2B_MODEL={TXGEMMA_2B}\n\nMEDIC_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2\nMEDIC_DATA_DIR=/kaggle/working/AMR-Guard/data\nMEDIC_CHROMA_DB_DIR=/kaggle/working/AMR-Guard/data/chroma_db\n\"\"\".strip()\n\nwith open(\"/kaggle/working/AMR-Guard/.env\", \"w\") as f:\n    f.write(env)\nprint(\".env written\")"
+   "source": [
+    "# Write .env\n",
+    "env = f\"\"\"\n",
+    "MEDIC_ENV=kaggle\n",
+    "MEDIC_QUANTIZATION=4bit\n",
+    "\n",
+    "# Agent 1, 2, 4 — MedGemma 4B IT\n",
+    "MEDIC_LOCAL_MEDGEMMA_4B_MODEL={MEDGEMMA_4B}\n",
+    "\n",
+    "# Agent 3 — MedGemma 27B Text IT  (subbed with 4B for Kaggle T4)\n",
+    "# To use full 27B: set to google/medgemma-27b-text-it\n",
+    "MEDIC_LOCAL_MEDGEMMA_27B_MODEL={MEDGEMMA_4B}\n",
+    "\n",
+    "# Agent 4 safety — TxGemma 9B  (subbed with 2B for Kaggle T4)\n",
+    "# To use full 9B: set to google/txgemma-9b-predict\n",
+    "MEDIC_LOCAL_TXGEMMA_9B_MODEL={TXGEMMA_2B}\n",
+    "MEDIC_LOCAL_TXGEMMA_2B_MODEL={TXGEMMA_2B}\n",
+    "\n",
+    "MEDIC_EMBEDDING_MODEL=sentence-transformers/all-MiniLM-L6-v2\n",
+    "MEDIC_DATA_DIR=/kaggle/working/AMR-Guard/data\n",
+    "MEDIC_CHROMA_DB_DIR=/kaggle/working/AMR-Guard/data/chroma_db\n",
+    "\"\"\".strip()\n",
+    "\n",
+    "with open(\"/kaggle/working/AMR-Guard/.env\", \"w\") as f:\n",
+    "    f.write(env)\n",
+    "print(\".env written\")"
+   ]
   },
   {
    "cell_type": "code",
@@ -467,6 +497,7 @@
     "    \"\"\"Stage 1 — Empirical: no lab results.\n",
     "    Active models: MedGemma 4B (Agent 1) → MedGemma 4B + TxGemma 2B (Agent 4).\n",
     "    \"\"\"\n",
+    "    logger.info(f\"Starting empirical scenario: {infection_site} infection\")\n",
     "    patient_data = _build_patient_data(\n",
     "        age, weight, height, sex, creatinine,\n",
     "        infection_site, suspected_source,\n",
@@ -474,8 +505,11 @@
     "    )\n",
     "    try:\n",
     "        from src.graph import run_pipeline\n",
+    "        logger.info(\"Calling run_pipeline...\")\n",
     "        result = run_pipeline(patient_data, labs_raw_text=None)\n",
+    "        logger.info(\"Pipeline completed successfully\")\n",
     "    except Exception as exc:\n",
+    "        logger.error(f\"Pipeline failed: {exc}\", exc_info=True)\n",
     "        result = _demo_result(patient_data, None)\n",
     "        result[\"errors\"].append(f\"[Demo mode — pipeline error: {exc}]\")\n",
     "    return format_recommendation(result)\n",
@@ -489,6 +523,7 @@
     "    Active models: MedGemma 4B (Agents 1, 2) → MedGemma 4B→27B sub (Agent 3)\n",
     "                   → MedGemma 4B + TxGemma 2B (Agent 4).\n",
     "    \"\"\"\n",
+    "    logger.info(f\"Starting targeted scenario: {infection_site} infection with lab data\")\n",
     "    patient_data = _build_patient_data(\n",
     "        age, weight, height, sex, creatinine,\n",
     "        infection_site, suspected_source,\n",
@@ -497,8 +532,11 @@
     "    labs = labs_text.strip() if labs_text else None\n",
     "    try:\n",
     "        from src.graph import run_pipeline\n",
+    "        logger.info(\"Calling run_pipeline with lab data...\")\n",
     "        result = run_pipeline(patient_data, labs_raw_text=labs)\n",
+    "        logger.info(\"Pipeline completed successfully\")\n",
     "    except Exception as exc:\n",
+    "        logger.error(f\"Pipeline failed: {exc}\", exc_info=True)\n",
     "        result = _demo_result(patient_data, labs)\n",
     "        result[\"errors\"].append(f\"[Demo mode — pipeline error: {exc}]\")\n",
     "    return format_recommendation(result), format_lab_analysis(result)\n",
@@ -507,6 +545,56 @@
     "print(\"Helper functions loaded.\")"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "57203904",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Enable logging to see what's happening during pipeline execution\n",
+    "import logging\n",
+    "logging.basicConfig(\n",
+    "    level=logging.INFO,\n",
+    "    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'\n",
+    ")\n",
+    "logger = logging.getLogger(__name__)\n",
+    "logger.info(\"Logging enabled for pipeline debugging\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "69c8a263",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Test model configuration and availability\n",
+    "print(\"Testing model configuration...\")\n",
+    "from src.config import get_settings\n",
+    "from src.loader import _get_model_path, _is_multimodal\n",
+    "\n",
+    "settings = get_settings()\n",
+    "print(f\"Environment: {settings.environment}\")\n",
+    "print(f\"Quantization: {settings.quantization}\")\n",
+    "print(f\"\\nConfigured models:\")\n",
+    "print(f\"  MedGemma 4B:  {settings.medgemma_4b_model}\")\n",
+    "print(f\"  MedGemma 27B: {settings.medgemma_27b_model}\")\n",
+    "print(f\"  TxGemma 2B:   {settings.txgemma_2b_model}\")\n",
+    "print(f\"  TxGemma 9B:   {settings.txgemma_9b_model}\")\n",
+    "\n",
+    "print(f\"\\nModel architectures:\")\n",
+    "for model_name in [\"medgemma_4b\", \"txgemma_2b\"]:\n",
+    "    try:\n",
+    "        path = _get_model_path(model_name)\n",
+    "        is_mm = _is_multimodal(path)\n",
+    "        print(f\"  {model_name}: {'multimodal' if is_mm else 'causal LM'}\")\n",
+    "    except Exception as e:\n",
+    "        print(f\"  {model_name}: ERROR - {e}\")\n",
+    "\n",
+    "print(\"\\n✓ Configuration validated\")"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -647,4 +735,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
+}

src/graph.py

@@ -91,16 +91,26 @@ def run_pipeline(patient_data: dict, labs_raw_text: str | None = None) -> Infect
         "country_or_region": patient_data.get("country_or_region"),
         "vitals": patient_data.get("vitals", {}),
         "stage": "targeted" if labs_raw_text else "empirical",
+        "errors": [],
+        "safety_warnings": [],
     }
 
     if labs_raw_text:
         initial_state["labs_raw_text"] = labs_raw_text
 
     logger.info(f"Starting pipeline (stage: {initial_state['stage']})")
-    compiled = build_infection_graph().compile()
-    final_state = compiled.invoke(initial_state)
-    logger.info("Pipeline complete")
-    return final_state
+    logger.info(f"Patient: {patient_data.get('age_years')}y, {patient_data.get('sex')}, infection: {patient_data.get('infection_site')}")
+    
+    try:
+        compiled = build_infection_graph().compile()
+        logger.info("Graph compiled successfully")
+        final_state = compiled.invoke(initial_state)
+        logger.info("Pipeline complete")
+        return final_state
+    except Exception as e:
+        logger.error(f"Pipeline execution failed: {e}", exc_info=True)
+        initial_state["errors"].append(f"Pipeline error: {str(e)}")
+        return initial_state
 
 
 def run_empirical_pipeline(patient_data: dict) -> InfectionState:

src/loader.py

@@ -9,13 +9,15 @@ logger = logging.getLogger(__name__)
 
 TextModelName = Literal["medgemma_4b", "medgemma_27b", "txgemma_9b", "txgemma_2b"]
 
+# MedGemma 4B IT is a vision-language model (Gemma3ForConditionalGeneration).
+# It must be loaded with AutoModelForImageTextToText + AutoProcessor.
+# All other models (medgemma-27b-text-it, txgemma-*) are causal LMs.
+# On Kaggle T4, medgemma_27b is substituted with medgemma-4b-it (also multimodal),
+# so we detect the architecture dynamically from the model config.
+_MULTIMODAL_ARCHITECTURES = {"Gemma3ForConditionalGeneration"}
 
-@lru_cache(maxsize=8)
-def _get_local_causal_lm(model_name: TextModelName):
-    """Load a local HuggingFace causal LM and return a generation callable."""
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    import torch
 
+def _get_model_path(model_name: TextModelName) -> str:
     settings = get_settings()
     model_path_map: Dict[TextModelName, Optional[str]] = {
         "medgemma_4b": settings.medgemma_4b_model,
@@ -23,21 +25,77 @@ def _get_local_causal_lm(model_name: TextModelName):
         "txgemma_9b": settings.txgemma_9b_model,
         "txgemma_2b": settings.txgemma_2b_model,
     }
-
     model_path = model_path_map[model_name]
     if not model_path:
         raise RuntimeError(
             f"No local model path configured for {model_name}. "
             f"Set MEDIC_LOCAL_*_MODEL in your environment or .env."
         )
+    return model_path
 
+
+def _get_load_kwargs() -> Dict[str, Any]:
+    settings = get_settings()
     load_kwargs: Dict[str, Any] = {"device_map": "auto"}
     if settings.quantization == "4bit":
         from transformers import BitsAndBytesConfig
         load_kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True)
+    return load_kwargs
+
+
+@lru_cache(maxsize=8)
+def _get_local_multimodal(model_name: TextModelName):
+    """Load a multimodal model (e.g. MedGemma 4B IT) and return a text generation callable."""
+    from transformers import AutoModelForImageTextToText, AutoProcessor
+    import torch
+
+    model_path = _get_model_path(model_name)
+    load_kwargs = _get_load_kwargs()
+
+    logger.info(f"Loading multimodal model: {model_path} with kwargs: {load_kwargs}")
+    processor = AutoProcessor.from_pretrained(model_path)
+    logger.info(f"Processor loaded for {model_path}")
+    model = AutoModelForImageTextToText.from_pretrained(model_path, **load_kwargs)
+    logger.info(f"Model loaded successfully: {model_path}")
+
+    def _call(prompt: str, max_new_tokens: int = 512, temperature: float = 0.2, **generate_kwargs: Any) -> str:
+        # Build a chat-style input for text-only queries
+        messages = [{"role": "user", "content": [{"type": "text", "text": prompt}]}]
+        inputs = processor.apply_chat_template(
+            messages, add_generation_prompt=True, tokenize=True,
+            return_dict=True, return_tensors="pt",
+        ).to(model.device)
+
+        do_sample = temperature > 0
+        with torch.no_grad():
+            output_ids = model.generate(
+                **inputs,
+                do_sample=do_sample,
+                temperature=temperature if do_sample else None,
+                max_new_tokens=max_new_tokens,
+                **generate_kwargs,
+            )
+        # Decode only the newly generated tokens
+        generated_ids = output_ids[0, inputs["input_ids"].shape[1]:]
+        return processor.decode(generated_ids, skip_special_tokens=True).strip()
+
+    return _call
 
+
+@lru_cache(maxsize=8)
+def _get_local_causal_lm(model_name: TextModelName):
+    """Load a causal LM (e.g. TxGemma, MedGemma 27B text) and return a generation callable."""
+    from transformers import AutoModelForCausalLM, AutoTokenizer
+    import torch
+
+    model_path = _get_model_path(model_name)
+    load_kwargs = _get_load_kwargs()
+
+    logger.info(f"Loading causal LM: {model_path} with kwargs: {load_kwargs}")
     tokenizer = AutoTokenizer.from_pretrained(model_path)
+    logger.info(f"Tokenizer loaded for {model_path}")
     model = AutoModelForCausalLM.from_pretrained(model_path, **load_kwargs)
+    logger.info(f"Model loaded successfully: {model_path}")
 
     def _call(prompt: str, max_new_tokens: int = 512, temperature: float = 0.2, **generate_kwargs: Any) -> str:
         inputs = {k: v.to(model.device) for k, v in tokenizer(prompt, return_tensors="pt").items()}
@@ -46,7 +104,7 @@ def _get_local_causal_lm(model_name: TextModelName):
             output_ids = model.generate(
                 **inputs,
                 do_sample=do_sample,
-                temperature=temperature if do_sample else 0.0,
+                temperature=temperature if do_sample else None,
                 max_new_tokens=max_new_tokens,
                 **generate_kwargs,
             )
@@ -57,11 +115,25 @@ def _get_local_causal_lm(model_name: TextModelName):
     return _call
 
 
+def _is_multimodal(model_path: str) -> bool:
+    """Check if a model uses a multimodal architecture by inspecting its config."""
+    from transformers import AutoConfig
+    try:
+        config = AutoConfig.from_pretrained(model_path)
+        architectures = getattr(config, "architectures", []) or []
+        return bool(set(architectures) & _MULTIMODAL_ARCHITECTURES)
+    except Exception:
+        return False
+
+
 @lru_cache(maxsize=32)
 def get_text_model(
     model_name: TextModelName = "medgemma_4b",
 ) -> Callable[..., str]:
     """Return a cached callable for the requested model."""
+    model_path = _get_model_path(model_name)
+    if _is_multimodal(model_path):
+        return _get_local_multimodal(model_name)
     return _get_local_causal_lm(model_name)
 
 
@@ -73,5 +145,13 @@ def run_inference(
     **kwargs: Any,
 ) -> str:
     """Run inference with the specified model. This is the primary entry point for agents."""
-    model = get_text_model(model_name=model_name)
-    return model(prompt, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
+    logger.info(f"Running inference with {model_name}, max_tokens={max_new_tokens}, temp={temperature}")
+    try:
+        model = get_text_model(model_name=model_name)
+        logger.info(f"Model {model_name} loaded successfully")
+        result = model(prompt, max_new_tokens=max_new_tokens, temperature=temperature, **kwargs)
+        logger.info(f"Inference complete, response length: {len(result)} chars")
+        return result
+    except Exception as e:
+        logger.error(f"Inference failed for {model_name}: {e}", exc_info=True)
+        raise