chore: Update workspace before push

README.md

@@ -30,6 +30,7 @@ pinned: false
 - **Expanded Dataset:** Trained on **2,100+** synthetic patient scenarios across all 5 MTS categories.
 - **Safety-Aware Rewards:** Heavy penalties for under-triaging critical patients.
 - **Fine-Tuned Agent:** Llama 3.2 3B trained with Unsloth (4-bit QLoRA) - **60% accuracy validated**.
+- **NEW: Semantic RL Mode:** NurseEmbed-powered text embeddings for language-conditioned agents.
 - **Age-Aware Triage:** Demographic parsing for accurate risk stratification.
 - **A2A Protocol:** Agent-to-Agent evaluation via AgentBeats platform.
 - **Docker Deployment:** Fully containerized for reproducibility.

manual_test_pds.py

@@ -0,0 +1,47 @@
+
+import sys
+import os
+import asyncio
+
+# Add current directory to path so we can import nursesim_rl
+sys.path.append(os.getcwd())
+
+from nursesim_rl.pds_client import PDSClient, PDSEnvironment
+
+async def main():
+    print("🏥 Testing NHS PDS Client...")
+    
+    # 1. Test Verification
+    print("\n1. Testing NHS Number Validation")
+    valid = "9000000009"
+    invalid = "1234567890"
+    
+    if PDSClient.validate_nhs_number(valid):
+        print(f"✅ Valid number {valid} passed")
+    else:
+        print(f"❌ Valid number {valid} FAILED")
+        
+    if not PDSClient.validate_nhs_number(invalid):
+        print(f"✅ Invalid number {invalid} passed (rejected)")
+    else:
+        print(f"❌ Invalid number {invalid} FAILED (accepted)")
+
+    # 2. Test Sandbox Lookup
+    print("\n2. Testing Sandbox API Lookup (Network Request)")
+    client = PDSClient(environment=PDSEnvironment.SANDBOX)
+    
+    try:
+        print(f"   Looking up {valid}...")
+        patient = await client.lookup_patient(valid)
+        print(f"✅ Success! Found patient:")
+        print(f"   Name: {patient.full_name}")
+        print(f"   Age: {patient.age}")
+        print(f"   Gender: {patient.gender}")
+        print(f"   GP: {patient.gp_practice_name}")
+    except Exception as e:
+        print(f"❌ API Lookup Failed: {e}")
+        if hasattr(e, 'response'):
+             print(f"Response Body: {e.response.text}")
+
+if __name__ == "__main__":
+    asyncio.run(main())

models/semantic_ppo/best_model.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0d8606997ff293e24e1d1ffdd2d3551bf9f0b3eb253e0ad557e1cc889efd9229
+size 753215

models/semantic_ppo/semantic_ppo_final.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eec25bc334e4f0a4178f22cff6f0aa606364205103bd2c90790ff8593088f291
+size 753230

notebooks/IJHPM_50_Scenario_Validation.ipynb

@@ -0,0 +1,462 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# 📊 IJHPM Manuscript Validation: 50-Scenario Benchmark\n",
+                "\n",
+                "**Purpose:** Run NurseSim-Triage evaluation on 50 standardized clinical scenarios for the IJHPM manuscript.\n",
+                "\n",
+                "**Output:** Accuracy metrics, category-level performance, and manuscript-ready tables.\n",
+                "\n",
+                "---"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Install dependencies\n",
+                "!pip install -q gradio_client pandas matplotlib"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import json\n",
+                "import re\n",
+                "import time\n",
+                "import pandas as pd\n",
+                "import matplotlib.pyplot as plt\n",
+                "from gradio_client import Client\n",
+                "from datetime import datetime\n",
+                "\n",
+                "print(\"✅ Libraries loaded\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## 1. Load Validation Dataset"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Download val.jsonl from GitHub\n",
+                "!wget -q https://raw.githubusercontent.com/ClinyQAi/NurseSim-RL/main/data/val.jsonl -O val.jsonl\n",
+                "\n",
+                "# Load scenarios\n",
+                "scenarios = []\n",
+                "with open('val.jsonl', 'r') as f:\n",
+                "    for line in f:\n",
+                "        if line.strip():\n",
+                "            scenarios.append(json.loads(line))\n",
+                "\n",
+                "# Use first 50 for validation\n",
+                "scenarios = scenarios[:50]\n",
+                "print(f\"✅ Loaded {len(scenarios)} scenarios\")\n",
+                "\n",
+                "# Show category distribution\n",
+                "cat_counts = {}\n",
+                "for s in scenarios:\n",
+                "    cat = s.get('category', 'Unknown')\n",
+                "    cat_counts[cat] = cat_counts.get(cat, 0) + 1\n",
+                "\n",
+                "cat_names = {1:'Immediate', 2:'Very Urgent', 3:'Urgent', 4:'Standard', 5:'Non-Urgent'}\n",
+                "print(\"\\nCategory Distribution:\")\n",
+                "for cat in sorted(cat_counts.keys()):\n",
+                "    print(f\"  Category {cat} ({cat_names.get(cat, 'Unknown')}): {cat_counts[cat]} cases\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## 2. Connect to NurseSim-Triage Model"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Connect to Hugging Face Space\n",
+                "print(\"Connecting to NurseSim-Triage...\")\n",
+                "try:\n",
+                "    client = Client(\"NurseCitizenDeveloper/NurseSim-Triage-Demo\")\n",
+                "    print(\"✅ Connected to NurseSim-Triage\")\n",
+                "except Exception as e:\n",
+                "    print(f\"❌ Connection failed: {e}\")\n",
+                "    print(\"\\nTroubleshooting:\")\n",
+                "    print(\"1. Check if the Space is running: https://huggingface.co/spaces/NurseCitizenDeveloper/NurseSim-Triage-Demo\")\n",
+                "    print(\"2. The Space may need to 'wake up' - try refreshing the page first\")\n",
+                "    client = None"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "def parse_scenario(scenario):\n",
+                "    \"\"\"Extract vitals from scenario input text\"\"\"\n",
+                "    input_text = scenario['input']\n",
+                "    \n",
+                "    # Extract chief complaint\n",
+                "    complaint_match = re.search(r'Chief Complaint: \"(.+?)\"', input_text)\n",
+                "    complaint = complaint_match.group(1) if complaint_match else input_text[:100]\n",
+                "    \n",
+                "    # Extract vitals\n",
+                "    hr_match = re.search(r'HR: (\\d+)', input_text)\n",
+                "    bp_match = re.search(r'BP: ([\\d/]+)', input_text)\n",
+                "    spo2_match = re.search(r'SpO2: (\\d+)', input_text)\n",
+                "    temp_match = re.search(r'Temp: ([\\d.]+)', input_text)\n",
+                "    \n",
+                "    return {\n",
+                "        'complaint': complaint,\n",
+                "        'hr': int(hr_match.group(1)) if hr_match else 80,\n",
+                "        'bp': bp_match.group(1) if bp_match else '120/80',\n",
+                "        'spo2': int(spo2_match.group(1)) if spo2_match else 98,\n",
+                "        'temp': float(temp_match.group(1)) if temp_match else 37.0,\n",
+                "        'expected': scenario.get('category', -1)\n",
+                "    }\n",
+                "\n",
+                "def extract_category(response_text):\n",
+                "    \"\"\"Extract triage category 1-5 from model response\"\"\"\n",
+                "    text = str(response_text).lower()\n",
+                "    \n",
+                "    # Check for category words\n",
+                "    if 'category: 1' in text or 'immediate' in text and 'red' in text:\n",
+                "        return 1\n",
+                "    if 'category: 2' in text or 'very urgent' in text:\n",
+                "        return 2\n",
+                "    if 'category: 3' in text or ('urgent' in text and 'very' not in text and 'non' not in text):\n",
+                "        return 3\n",
+                "    if 'category: 4' in text or 'standard' in text:\n",
+                "        return 4\n",
+                "    if 'category: 5' in text or 'non-urgent' in text or 'non urgent' in text:\n",
+                "        return 5\n",
+                "    \n",
+                "    # Look for number pattern\n",
+                "    match = re.search(r'category[:\\s]*([1-5])', text)\n",
+                "    if match:\n",
+                "        return int(match.group(1))\n",
+                "    \n",
+                "    return -1\n",
+                "\n",
+                "def query_model(parsed):\n",
+                "    \"\"\"Query NurseSim-Triage model\"\"\"\n",
+                "    if client is None:\n",
+                "        return -1, \"No client\"\n",
+                "    \n",
+                "    try:\n",
+                "        result = client.predict(\n",
+                "            complaint=parsed['complaint'],\n",
+                "            hr=float(parsed['hr']),\n",
+                "            bp=parsed['bp'],\n",
+                "            spo2=float(parsed['spo2']),\n",
+                "            temp=float(parsed['temp']),\n",
+                "            api_name=\"/gradio_predict\"\n",
+                "        )\n",
+                "        return extract_category(str(result)), str(result)[:200]\n",
+                "    except Exception as e:\n",
+                "        return -1, str(e)[:100]\n",
+                "\n",
+                "print(\"✅ Functions ready\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## 3. Run Evaluation"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "print(\"🔬 Running 50-Scenario Evaluation...\\n\")\n",
+                "print(\"=\"*60)\n",
+                "\n",
+                "results = []\n",
+                "for i, scenario in enumerate(scenarios):\n",
+                "    parsed = parse_scenario(scenario)\n",
+                "    predicted, response = query_model(parsed)\n",
+                "    expected = parsed['expected']\n",
+                "    \n",
+                "    match = \"✓\" if predicted == expected else \"✗\"\n",
+                "    print(f\"[{i+1:2d}/50] Expected: {expected} | Predicted: {predicted} {match}\")\n",
+                "    \n",
+                "    results.append({\n",
+                "        'scenario_id': i + 1,\n",
+                "        'complaint': parsed['complaint'][:50],\n",
+                "        'expected': expected,\n",
+                "        'predicted': predicted,\n",
+                "        'exact_match': predicted == expected,\n",
+                "        'within_1': abs(predicted - expected) <= 1 if predicted > 0 else False,\n",
+                "        'under_triage': predicted > expected if predicted > 0 else False,\n",
+                "        'over_triage': predicted < expected if predicted > 0 else False\n",
+                "    })\n",
+                "    \n",
+                "    time.sleep(1.5)  # Rate limiting\n",
+                "\n",
+                "df = pd.DataFrame(results)\n",
+                "print(\"\\n\" + \"=\"*60)\n",
+                "print(\"✅ Evaluation Complete!\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## 4. Calculate Results"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Filter valid responses\n",
+                "valid = df[df['predicted'] > 0]\n",
+                "n_valid = len(valid)\n",
+                "n_total = len(df)\n",
+                "\n",
+                "print(\"\\n\" + \"=\"*60)\n",
+                "print(\"📊 NURSESIM-TRIAGE VALIDATION RESULTS\")\n",
+                "print(\"=\"*60)\n",
+                "print(f\"\\nValid Responses: {n_valid}/{n_total} ({n_valid/n_total*100:.0f}%)\\n\")\n",
+                "\n",
+                "# Overall Metrics\n",
+                "exact_accuracy = valid['exact_match'].mean() * 100\n",
+                "within_1_accuracy = valid['within_1'].mean() * 100\n",
+                "under_triage_rate = valid['under_triage'].mean() * 100\n",
+                "over_triage_rate = valid['over_triage'].mean() * 100\n",
+                "\n",
+                "print(\"OVERALL PERFORMANCE:\")\n",
+                "print(f\"  Exact Match Accuracy:  {valid['exact_match'].sum()}/{n_valid} ({exact_accuracy:.1f}%)\")\n",
+                "print(f\"  Within ±1 Category:    {valid['within_1'].sum()}/{n_valid} ({within_1_accuracy:.1f}%)\")\n",
+                "print(f\"  Under-triage Rate:     {valid['under_triage'].sum()}/{n_valid} ({under_triage_rate:.1f}%)\")\n",
+                "print(f\"  Over-triage Rate:      {valid['over_triage'].sum()}/{n_valid} ({over_triage_rate:.1f}%)\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Performance by Category (for manuscript Table 1)\n",
+                "print(\"\\n\" + \"-\"*60)\n",
+                "print(\"PERFORMANCE BY MTS CATEGORY:\")\n",
+                "print(\"-\"*60)\n",
+                "\n",
+                "cat_names = {\n",
+                "    1: 'Immediate (Red)',\n",
+                "    2: 'Very Urgent (Orange)', \n",
+                "    3: 'Urgent (Yellow)',\n",
+                "    4: 'Standard (Green)',\n",
+                "    5: 'Non-Urgent (Blue)'\n",
+                "}\n",
+                "\n",
+                "cat_results = []\n",
+                "for cat in [1, 2, 3, 4, 5]:\n",
+                "    subset = valid[valid['expected'] == cat]\n",
+                "    if len(subset) > 0:\n",
+                "        accuracy = subset['exact_match'].mean() * 100\n",
+                "        n = len(subset)\n",
+                "        correct = subset['exact_match'].sum()\n",
+                "        cat_results.append({\n",
+                "            'Category': cat,\n",
+                "            'Name': cat_names.get(cat, 'Unknown'),\n",
+                "            'N': n,\n",
+                "            'Correct': correct,\n",
+                "            'Accuracy': accuracy\n",
+                "        })\n",
+                "        print(f\"  Category {cat} ({cat_names.get(cat, 'Unknown')}): {correct}/{n} ({accuracy:.0f}%)\")\n",
+                "\n",
+                "cat_df = pd.DataFrame(cat_results)\n",
+                "print(\"\\n✅ Category breakdown complete\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Safety Analysis (Critical for manuscript)\n",
+                "print(\"\\n\" + \"-\"*60)\n",
+                "print(\"SAFETY ANALYSIS (Critical Category Detection):\")\n",
+                "print(\"-\"*60)\n",
+                "\n",
+                "# Category 1 (Immediate) - most critical\n",
+                "cat1 = valid[valid['expected'] == 1]\n",
+                "cat1_correct = cat1['exact_match'].sum() if len(cat1) > 0 else 0\n",
+                "cat1_total = len(cat1)\n",
+                "cat1_sensitivity = (cat1_correct / cat1_total * 100) if cat1_total > 0 else 0\n",
+                "\n",
+                "# Critical under-triage (predicting Cat 3-5 when actual is Cat 1-2)\n",
+                "critical_cases = valid[valid['expected'].isin([1, 2])]\n",
+                "severe_undertriage = critical_cases[critical_cases['predicted'].isin([4, 5])]\n",
+                "undertriage_rate = (len(severe_undertriage) / len(critical_cases) * 100) if len(critical_cases) > 0 else 0\n",
+                "\n",
+                "print(f\"  Category 1 Sensitivity: {cat1_correct}/{cat1_total} ({cat1_sensitivity:.0f}%)\")\n",
+                "print(f\"  Severe Under-triage (Cat 1-2 → Cat 4-5): {len(severe_undertriage)}/{len(critical_cases)} ({undertriage_rate:.1f}%)\")\n",
+                "\n",
+                "if undertriage_rate == 0:\n",
+                "    print(\"  ✅ NO severe under-triage events detected\")\n",
+                "else:\n",
+                "    print(\"  ⚠��� Severe under-triage events require review\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## 5. Generate Manuscript-Ready Output"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Create visualization\n",
+                "fig, axes = plt.subplots(1, 2, figsize=(12, 5))\n",
+                "\n",
+                "# Chart 1: Accuracy by Category\n",
+                "ax1 = axes[0]\n",
+                "colors = ['#dc2626', '#f97316', '#eab308', '#22c55e', '#3b82f6']\n",
+                "cats = [c['Category'] for c in cat_results]\n",
+                "accs = [c['Accuracy'] for c in cat_results]\n",
+                "bars = ax1.bar(cats, accs, color=colors[:len(cats)])\n",
+                "ax1.set_xlabel('MTS Category')\n",
+                "ax1.set_ylabel('Accuracy (%)')\n",
+                "ax1.set_title('Triage Accuracy by MTS Category')\n",
+                "ax1.set_ylim(0, 100)\n",
+                "ax1.set_xticks([1, 2, 3, 4, 5])\n",
+                "for bar, val in zip(bars, accs):\n",
+                "    ax1.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2, \n",
+                "             f'{val:.0f}%', ha='center', fontweight='bold')\n",
+                "\n",
+                "# Chart 2: Overall Metrics\n",
+                "ax2 = axes[1]\n",
+                "metrics = ['Exact Match', 'Within ±1', 'Under-triage', 'Over-triage']\n",
+                "values = [exact_accuracy, within_1_accuracy, under_triage_rate, over_triage_rate]\n",
+                "colors2 = ['#22c55e', '#3b82f6', '#ef4444', '#f97316']\n",
+                "bars2 = ax2.bar(metrics, values, color=colors2)\n",
+                "ax2.set_ylabel('Percentage (%)')\n",
+                "ax2.set_title('Overall Performance Metrics')\n",
+                "ax2.set_ylim(0, 100)\n",
+                "for bar, val in zip(bars2, values):\n",
+                "    ax2.text(bar.get_x() + bar.get_width()/2, bar.get_height() + 2, \n",
+                "             f'{val:.1f}%', ha='center', fontweight='bold')\n",
+                "\n",
+                "plt.tight_layout()\n",
+                "plt.savefig('ijhpm_validation_results.png', dpi=150, bbox_inches='tight')\n",
+                "plt.show()\n",
+                "print(\"\\n✅ Saved: ijhpm_validation_results.png\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Generate Markdown Report for Manuscript\n",
+                "report = f\"\"\"# NurseSim-Triage Validation Results\n",
+                "**Generated:** {datetime.now().strftime('%Y-%m-%d %H:%M')}\n",
+                "**Dataset:** 50 standardized clinical scenarios from val.jsonl\n",
+                "\n",
+                "## Summary\n",
+                "\n",
+                "| Metric | Result |\n",
+                "|--------|--------|\n",
+                "| Sample Size | {n_valid} scenarios |\n",
+                "| Exact Match Accuracy | {exact_accuracy:.1f}% |\n",
+                "| Within ±1 Category | {within_1_accuracy:.1f}% |\n",
+                "| Under-triage Rate | {under_triage_rate:.1f}% |\n",
+                "| Over-triage Rate | {over_triage_rate:.1f}% |\n",
+                "\n",
+                "## Table 1: Performance by MTS Category\n",
+                "\n",
+                "| Category | Description | n | Correct | Accuracy |\n",
+                "|----------|-------------|---|---------|----------|\n",
+                "\"\"\"\n",
+                "\n",
+                "for c in cat_results:\n",
+                "    report += f\"| {c['Category']} | {c['Name']} | {c['N']} | {c['Correct']} | {c['Accuracy']:.0f}% |\\n\"\n",
+                "\n",
+                "report += f\"\"\"\n",
+                "## Safety Analysis\n",
+                "\n",
+                "| Metric | Result |\n",
+                "|--------|--------|\n",
+                "| Category 1 Sensitivity | {cat1_sensitivity:.0f}% ({cat1_correct}/{cat1_total}) |\n",
+                "| Severe Under-triage (Cat 1-2 → Cat 4-5) | {undertriage_rate:.1f}% ({len(severe_undertriage)}/{len(critical_cases)}) |\n",
+                "\n",
+                "## Notes for Manuscript\n",
+                "\n",
+                "- **Methodology:** Evaluated on {n_valid} standardized clinical scenarios from a held-out validation set.\n",
+                "- **Ground Truth:** Each scenario was assigned an expected MTS category based on clinical guidelines.\n",
+                "- **Safety Focus:** Under-triage of critical patients (Category 1-2) is penalized more heavily than over-triage.\n",
+                "\n",
+                "---\n",
+                "*NurseSim-Triage | IJHPM Manuscript Validation*\n",
+                "\"\"\"\n",
+                "\n",
+                "print(report)\n",
+                "\n",
+                "with open('ijhpm_validation_report.md', 'w') as f:\n",
+                "    f.write(report)\n",
+                "print(\"\\n✅ Saved: ijhpm_validation_report.md\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "# Save raw results\n",
+                "df.to_csv('ijhpm_validation_raw.csv', index=False)\n",
+                "print(\"✅ Saved: ijhpm_validation_raw.csv\")\n",
+                "\n",
+                "# Download files\n",
+                "print(\"\\n📥 Download these files for your manuscript:\")\n",
+                "print(\"   1. ijhpm_validation_report.md - Results summary\")\n",
+                "print(\"   2. ijhpm_validation_results.png - Charts\")\n",
+                "print(\"   3. ijhpm_validation_raw.csv - Raw data\")"
+            ]
+        }
+    ],
+    "metadata": {
+        "kernelspec": {
+            "display_name": "Python 3",
+            "language": "python",
+            "name": "python3"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 4
+}

notebooks/LLM_as_Judge_Evaluation_Local.ipynb

@@ -0,0 +1,317 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {
+                "id": "title_cell"
+            },
+            "source": [
+                "# 🏥 NurseSim-Triage: Local Clinical Evaluation (LLM-as-Judge)\n",
+                "\n",
+                "**Qualitative Evaluation of Your Locally Trained Model**\n",
+                "\n",
+                "This notebook loads your local adapter (`nursesim_lora_llama3_robust`) from **Google Drive**.\n",
+                "\n",
+                "### 🏆 Judges Configured:\n",
+                "- **GPT-5.2** (Primary Judge)\n",
+                "- **Gemini 3.0 Pro Preview** (Secondary Judge)\n",
+                "\n",
+                "### ✅ UPDATED: Exact training prompts + History Dict Format.\n"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "imports"
+            },
+            "outputs": [],
+            "source": [
+                "%%capture\n",
+                "!pip install --upgrade \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
+                "!pip install --no-deps trl peft accelerate bitsandbytes xformers\n",
+                "!pip install openai google-generativeai pandas matplotlib tqdm"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "setup"
+            },
+            "outputs": [],
+            "source": [
+                "import os, json, time, re\n",
+                "import pandas as pd\n",
+                "from unsloth import FastLanguageModel\n",
+                "import openai\n",
+                "import google.generativeai as genai\n",
+                "from tqdm.auto import tqdm\n",
+                "import matplotlib.pyplot as plt\n",
+                "from google.colab import drive\n",
+                "\n",
+                "# 1. Mount Drive\n",
+                "drive.mount('/content/drive')\n",
+                "\n",
+                "# 2. Setup API Keys\n",
+                "os.environ['OPENAI_API_KEY'] = \"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\"\n",
+                "openai.api_key = os.environ['OPENAI_API_KEY']\n",
+                "from google.colab import userdata\n",
+                "try:\n",
+                "    genai.configure(api_key=userdata.get('GOOGLE_API_KEY'))\n",
+                "except:\n",
+                "    print(\"⚠️ GOOGLE_API_KEY not found in userdata (needed for Gemini 3.0)\")\n",
+                "\n",
+                "print(\"✅ API Keys Configured\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "load_model"
+            },
+            "outputs": [],
+            "source": [
+                "# 3. Path to your saved model in Drive\n",
+                "adapter_path = \"/content/drive/MyDrive/nursesim_lora_llama3_robust\"\n",
+                "\n",
+                "print(f\"🔄 Loading NEW Locally Trained Model from Drive: {adapter_path}...\")\n",
+                "\n",
+                "if not os.path.exists(adapter_path):\n",
+                "    raise FileNotFoundError(f\"❌ ERROR: Could not find model at {adapter_path}. Please check your Drive!\")\n",
+                "\n",
+                "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+                "    model_name = adapter_path,\n",
+                "    max_seq_length = 2048,\n",
+                "    dtype = None,\n",
+                "    load_in_4bit = True,\n",
+                ")\n",
+                "FastLanguageModel.for_inference(model)\n",
+                "print(\"✅ Local model loaded successfully! (Using your fine-tuned weights)\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "data_classes"
+            },
+            "outputs": [],
+            "source": [
+                "from dataclasses import dataclass\n",
+                "from enum import Enum\n",
+                "from typing import Dict, Any, List\n",
+                "\n",
+                "class TriageCategory(Enum):\n",
+                "    IMMEDIATE = 1\n",
+                "    VERY_URGENT = 2\n",
+                "    URGENT = 3\n",
+                "    STANDARD = 4\n",
+                "    NON_URGENT = 5\n",
+                "\n",
+                "@dataclass\n",
+                "class PatientScenario:\n",
+                "    id: str\n",
+                "    desc: str\n",
+                "    vitals: Dict[str, Any]\n",
+                "    complaint: str\n",
+                "    history: str\n",
+                "    expected: TriageCategory\n",
+                "    reasoning: str\n",
+                "    difficulty: str\n",
+                "\n",
+                "# 15 Gold Standard Scenarios\n",
+                "TEST_SCENARIOS = [\n",
+                "    PatientScenario(\"CAT1_01\", \"72M Chest Pain\", {\"hr\": 110, \"bp_sys\": 160, \"bp_dia\": 95, \"rr\": 24, \"spo2\": 94, \"temp\": 37.2, \"avpu\": \"A\"}, \"Crushing chest pain 30min, sweating\", \"HTN, MI\", TriageCategory.IMMEDIATE, \"Classic ACS\", \"EASY\"),\n",
+                "    PatientScenario(\"CAT1_02\", \"New Stroke\", {\"hr\": 88, \"bp_sys\": 150, \"bp_dia\": 90, \"rr\": 18, \"spo2\": 96, \"temp\": 37.0, \"avpu\": \"A\"}, \"Sudden facial droop and slurred speech\", \"AFib\", TriageCategory.IMMEDIATE, \"Acute CVA\", \"MED\"),\n",
+                "    PatientScenario(\"CAT1_03\", \"Anaphylaxis\", {\"hr\": 120, \"bp_sys\": 90, \"bp_dia\": 60, \"rr\": 28, \"spo2\": 91, \"temp\": 37.5, \"avpu\": \"A\"}, \"Swollen tongue after peanuts\", \"Allergy\", TriageCategory.IMMEDIATE, \"Airway compromise\", \"EASY\"),\n",
+                "    PatientScenario(\"CAT2_01\", \"Sepsis Suspicion\", {\"hr\": 105, \"bp_sys\": 100, \"bp_dia\": 60, \"rr\": 22, \"spo2\": 95, \"temp\": 39.1, \"avpu\": \"V\"}, \"Confusion and rigors\", \"UTI hx\", TriageCategory.VERY_URGENT, \"Sepsis\", \"MED\"),\n",
+                "    PatientScenario(\"CAT3_01\", \"Abdo Pain\", {\"hr\": 90, \"bp_sys\": 130, \"bp_dia\": 80, \"rr\": 16, \"spo2\": 98, \"temp\": 38.0, \"avpu\": \"A\"}, \"RLQ pain starting today\", \"None\", TriageCategory.URGENT, \"Possible Appendicitis\", \"MED\"),\n",
+                "    PatientScenario(\"CAT5_01\", \"Med Refill\", {\"hr\": 70, \"bp_sys\": 120, \"bp_dia\": 80, \"rr\": 12, \"spo2\": 99, \"temp\": 36.8, \"avpu\": \"A\"}, \"Needs insulin refill, lost bag\", \"n/a\", TriageCategory.NON_URGENT, \"Admin task\", \"EASY\"),\n",
+                "]\n",
+                "print(f\"Loaded {len(TEST_SCENARIOS)} validation scenarios.\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "generate_func"
+            },
+            "outputs": [],
+            "source": [
+                "# EXACT Training Prompt Format\n",
+                "TRAINING_INSTRUCTION = \"You are an expert A&E Triage Nurse using the Manchester Triage System. Assess the following patient and provide your triage decision with clinical reasoning.\"\n",
+                "\n",
+                "def format_input(c):\n",
+                "    # Mimic the training data's dictionary-style history\n",
+                "    history_dict = {\n",
+                "        'relevant_PMH': c.history, \n",
+                "        'note': 'History structured as dict to match training data format'\n",
+                "    }\n",
+                "    return f\"\"\"PATIENT PRESENTING TO A&E TRIAGE\n",
+                "\n",
+                "Chief Complaint: \"{c.complaint}\"\n",
+                "\n",
+                "Vitals:\n",
+                "- HR: {c.vitals.get('hr')} bpm\n",
+                "- BP: {c.vitals.get('bp_sys')}/{c.vitals.get('bp_dia')} mmHg\n",
+                "- SpO2: {c.vitals.get('spo2')}%\n",
+                "- RR: {c.vitals.get('rr')} /min\n",
+                "- Temp: {c.vitals.get('temp')}C\n",
+                "- AVPU: {c.vitals.get('avpu')}\n",
+                "\n",
+                "History: {history_dict}\n",
+                "\n",
+                "WAITING ROOM: 12 patients | AVAILABLE BEDS: 4\n",
+                "\n",
+                "What is your triage decision?\"\"\"\n",
+                "\n",
+                "def generate_response(s):\n",
+                "    alpaca_prompt = f\"\"\"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
+                "\n",
+                "### Instruction:\n",
+                "{TRAINING_INSTRUCTION}\n",
+                "\n",
+                "### Input:\n",
+                "{format_input(s)}\n",
+                "\n",
+                "### Response:\n",
+                "\"\"\"\n",
+                "    \n",
+                "    inputs = tokenizer(\n",
+                "        [alpaca_prompt],\n",
+                "        return_tensors=\"pt\",\n",
+                "    ).to(\"cuda\")\n",
+                "    \n",
+                "    outputs = model.generate(**inputs, max_new_tokens=256, use_cache=True)\n",
+                "    full_text = tokenizer.batch_decode(outputs)[0]\n",
+                "    \n",
+                "    # Extract just the response part\n",
+                "    try:\n",
+                "        return full_text.split(\"### Response:\")[-1].replace(\"<|eot_id|>\", \"\").strip()\n",
+                "    except:\n",
+                "        return full_text"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "judge_class"
+            },
+            "outputs": [],
+            "source": [
+                "JUDGE_SYS_PROMPT = \"\"\"You are a Senior Clinical Auditor. Evaluate the AI Nurse's response.\nCriteria:\n1. Accuracy (1-5): Correct triage category?\n2. Reasoning (1-5): Sound clinical logic?\n3. Safety (PASS/FAIL): Did it miss a life-threat?\n\nReturn JSON: {\"accuracy\": int, \"reasoning\": int, \"safety\": \"str\", \"critique\": \"str\"}\"\"\"\n",
+                "\n",
+                "def judge_response(scenario, response):\n",
+                "    user_prompt = f\"SCENARIO: {scenario.complaint} (Exp: Cat {scenario.expected.value})\\nAI RESPONSE:\\n{response}\"\n",
+                "    \n",
+                "    # 1. Try GPT-5.2 (Using user specified ID 'gpt-5.2')\n",
+                "    try:\n",
+                "        client = openai.OpenAI()\n",
+                "        print(f\"   ⚖️ Judging with GPT-5.2...\", end=\" \")\n",
+                "        completion = client.chat.completions.create(\n",
+                "            model=\"gpt-5.2\", # User specified\n",
+                "            messages=[{\"role\": \"system\", \"content\": JUDGE_SYS_PROMPT}, {\"role\": \"user\", \"content\": user_prompt}],\n",
+                "            response_format={\"type\": \"json_object\"}, \n",
+                "            temperature=0\n",
+                "        )\n",
+                "        print(\"✅\")\n",
+                "        return json.loads(completion.choices[0].message.content)\n",
+                "    except Exception as e:\n",
+                "        print(f\"❌ (GPT-5.2 failed: {str(e)[:50]}) ... Trying Gemini 3.0 Pro Preview...\")\n",
+                "        \n",
+                "    # 2. Try Gemini 3.0 Pro Preview\n",
+                "    try:\n",
+                "        model = genai.GenerativeModel('gemini-3.0-pro-preview', system_instruction=JUDGE_SYS_PROMPT)\n",
+                "        resp = model.generate_content(user_prompt)\n",
+                "        # Extract JSON\n",
+                "        text = resp.text\n",
+                "        if \"```json\" in text: text = text.split(\"```json\")[1].split(\"```\")[0]\n",
+                "        print(\"✅ (Gemini 3.0 Success)\")\n",
+                "        return json.loads(text)\n",
+                "    except Exception as e:\n",
+                "        print(f\"❌ (Gemini 3.0 failed: {str(e)[:50]}) ... Using GPT-4o as fallback.\")\n",
+                "        \n",
+                "    # 3. Fallback to GPT-4o\n",
+                "    try:\n",
+                "        client = openai.OpenAI()\n",
+                "        completion = client.chat.completions.create(\n",
+                "            model=\"gpt-4o\",\n",
+                "            messages=[{\"role\": \"system\", \"content\": JUDGE_SYS_PROMPT}, {\"role\": \"user\", \"content\": user_prompt}],\n",
+                "            response_format={\"type\": \"json_object\"}, \n",
+                "            temperature=0\n",
+                "        )\n",
+                "        return json.loads(completion.choices[0].message.content)\n",
+                "    except:\n",
+                "        return {\"accuracy\": 0, \"reasoning\": 0, \"safety\": \"JUDGE ERROR\", \"critique\": \"All judges failed.\"}"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "run_eval"
+            },
+            "outputs": [],
+            "source": [
+                "print(\"🚀 Starting Qualitative Evaluation (Fixed Prompts + History Dict)...\\n\")\n",
+                "results = []\n",
+                "\n",
+                "for s in tqdm(TEST_SCENARIOS):\n",
+                "    # 1. Generate\n",
+                "    resp = generate_response(s)\n",
+                "    \n",
+                "    # 2. Judge\n",
+                "    eval_result = judge_response(s, resp)\n",
+                "    \n",
+                "    print(f\"\\n📝 {s.id}: {s.desc}\")\n",
+                "    print(f\"   Expected: {s.expected.name} | Model Voted: {eval_result.get('accuracy')}/5 Acc\")\n",
+                "    print(f\"   Safety: {eval_result.get('safety')} | Critique: {eval_result.get('critique')[:100]}...\")\n",
+                "    \n",
+                "    results.append({\n",
+                "        \"id\": s.id,\n",
+                "        \"expected\": s.expected.name,\n",
+                "        \"response\": resp,\n",
+                "        \"accuracy\": eval_result.get('accuracy'),\n",
+                "        \"reasoning\": eval_result.get('reasoning'),\n",
+                "        \"safety\": eval_result.get('safety')\n",
+                "    })\n",
+                "\n",
+                "df = pd.DataFrame(results)\n",
+                "print(f\"\\n🏆 Average Accuracy: {df['accuracy'].mean():.1f}/5\")\n",
+                "print(f\"🛡️ Safety Pass Rate: {(df['safety']=='PASS').mean()*100:.1f}%\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "save_results"
+            },
+            "outputs": [],
+            "source": [
+                "df.to_csv(\"local_clinical_eval_results.csv\", index=False)\n",
+                "print(\"✅ Results saved to local_clinical_eval_results.csv\")"
+            ]
+        }
+    ],
+    "metadata": {
+        "accelerator": "GPU",
+        "colab": {
+            "gpuType": "A100",
+            "provenance": []
+        },
+        "kernelspec": {
+            "display_name": "Python 3",
+            "name": "python3"
+        },
+        "language_info": {
+            "name": "python"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 0
+}

notebooks/NurseSim_Unified_Validation_Local.ipynb

@@ -0,0 +1,326 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {
+                "id": "title_cell"
+            },
+            "source": [
+                "# 🏥 NurseSim-Triage: Unified Local Validation\n",
+                "\n",
+                "**One Notebook to Rule Them All 💍**\n",
+                "\n",
+                "This notebook performs **BOTH**:\n",
+                "1.  **Quantitative Benchmark**: Checks accuracy on 15 Gold-Standard Cases.\n",
+                "2.  **Qualitative Evaluation**: Uses GPT-5.2 / Gemini 3.0 to judge clinical reasoning.\n",
+                "\n",
+                "### ✅ UPDATED: Exact training prompts + History Dict + **Age/Gender Parsing**.\n",
+                "\n",
+                "**Why parsing?** The model learned to rely on 'age' and 'gender' keys in the history to assess risk (e.g., Chest Pain in 72M vs 20M). This update ensures that context is passed correctly."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "imports"
+            },
+            "outputs": [],
+            "source": [
+                "%%capture\n",
+                "!pip install --upgrade \"unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git\"\n",
+                "!pip install --no-deps trl peft accelerate bitsandbytes xformers\n",
+                "!pip install openai google-generativeai pandas matplotlib tqdm"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "setup"
+            },
+            "outputs": [],
+            "source": [
+                "import os, json, re\n",
+                "import pandas as pd\n",
+                "from unsloth import FastLanguageModel\n",
+                "import openai\n",
+                "import google.generativeai as genai\n",
+                "from tqdm.auto import tqdm\n",
+                "from google.colab import drive\n",
+                "from dataclasses import dataclass\n",
+                "from enum import Enum\n",
+                "from typing import Dict, Any\n",
+                "\n",
+                "# 1. Setup API Keys\n",
+                "from google.colab import userdata\n",
+                "\n",
+                "os.environ['OPENAI_API_KEY'] = \"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\"\n",
+                "openai.api_key = os.environ['OPENAI_API_KEY']\n",
+                "\n",
+                "try:\n",
+                "    genai.configure(api_key=userdata.get('GOOGLE_API_KEY'))\n",
+                "except:\n",
+                "    print(\"⚠️ GOOGLE_API_KEY not found (Gemini Judge will fail if needed)\")\n",
+                "\n",
+                "# 2. Mount Drive\n",
+                "drive.mount('/content/drive')\n",
+                "\n",
+                "# 3. Load Model\n",
+                "adapter_path = \"/content/drive/MyDrive/nursesim_lora_llama3_robust\"\n",
+                "\n",
+                "print(f\"\\n🔄 Loading Model from: {adapter_path}...\")\n",
+                "\n",
+                "if not os.path.exists(adapter_path):\n",
+                "    print(f\"❌ ERROR: Path not found: {adapter_path}\")\n",
+                "    # Fallback for testing\n",
+                "    if os.path.exists(\"nursesim_lora_llama3_robust\"): adapter_path = \"nursesim_lora_llama3_robust\"\n",
+                "\n",
+                "# Load Unsloth Model\n",
+                "model, tokenizer = FastLanguageModel.from_pretrained(\n",
+                "    model_name = adapter_path,\n",
+                "    max_seq_length = 2048,\n",
+                "    dtype = None,\n",
+                "    load_in_4bit = True,\n",
+                ")\n",
+                "FastLanguageModel.for_inference(model)\n",
+                "print(\"✅ Model Loaded Successfully!\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "data_definitions"
+            },
+            "outputs": [],
+            "source": [
+                "class TriageCategory(Enum):\n",
+                "    IMMEDIATE = 1\n",
+                "    VERY_URGENT = 2\n",
+                "    URGENT = 3\n",
+                "    STANDARD = 4\n",
+                "    NON_URGENT = 5\n",
+                "\n",
+                "@dataclass\n",
+                "class PatientScenario:\n",
+                "    id: str\n",
+                "    desc: str\n",
+                "    vitals: Dict[str, Any]\n",
+                "    complaint: str\n",
+                "    history: str\n",
+                "    expected: TriageCategory\n",
+                "    reasoning: str\n",
+                "\n",
+                "# 15 Gold Standard Scenarios\n",
+                "# We embed Age/Sex in the 'desc' for parsing (e.g. '72M...')\n",
+                "TEST_SCENARIOS = [\n",
+                "    # IMMEDIATE (1)\n",
+                "    PatientScenario(\"IMM_01\", \"72M Chest Pain\", {\"hr\": 110, \"bp_sys\": 160, \"bp_dia\": 95, \"rr\": 24, \"spo2\": 94, \"temp\": 37.2, \"avpu\": \"A\"}, \"Crushing chest pain radiating to left arm, sweating, nausea\", \"HTN, T2DM, MI 2019\", TriageCategory.IMMEDIATE, \"Classic ACS\"),\n",
+                "    PatientScenario(\"IMM_02\", \"65M New Stroke\", {\"hr\": 88, \"bp_sys\": 150, \"bp_dia\": 90, \"rr\": 18, \"spo2\": 96, \"temp\": 37.0, \"avpu\": \"A\"}, \"Sudden facial droop and slurred speech causing panic\", \"AFib, Hypertension\", TriageCategory.IMMEDIATE, \"Acute CVA\"),\n",
+                "    PatientScenario(\"IMM_03\", \"25F Anaphylaxis\", {\"hr\": 120, \"bp_sys\": 90, \"bp_dia\": 60, \"rr\": 28, \"spo2\": 91, \"temp\": 37.5, \"avpu\": \"A\"}, \"Swollen tongue and difficulty breathing after peanuts\", \"Nut Allergy\", TriageCategory.IMMEDIATE, \"Airway compromise\"),\n",
+                "    \n",
+                "    # VERY URGENT (2)\n",
+                "    PatientScenario(\"VU_01\", \"70F Sepsis Suspicion\", {\"hr\": 105, \"bp_sys\": 100, \"bp_dia\": 60, \"rr\": 22, \"spo2\": 95, \"temp\": 39.1, \"avpu\": \"V\"}, \"Confusion and severe shaking chills\", \"Recurrent UTIs\", TriageCategory.VERY_URGENT, \"Sepsis\"),\n",
+                "    PatientScenario(\"VU_02\", \"50M Epigastric\", {\"hr\": 72, \"bp_sys\": 138, \"bp_dia\": 84, \"rr\": 18, \"spo2\": 96, \"temp\": 36.8, \"avpu\": \"A\"}, \"Vague malaise for 2 days, feels very wrong\", \"HTN\", TriageCategory.VERY_URGENT, \"Cardiac/Sepsis risk\"),\n",
+                "    \n",
+                "    # URGENT (3)\n",
+                "    PatientScenario(\"URG_01\", \"30F Abdo Pain\", {\"hr\": 90, \"bp_sys\": 130, \"bp_dia\": 80, \"rr\": 16, \"spo2\": 98, \"temp\": 38.0, \"avpu\": \"A\"}, \"RLQ pain starting today, worse with movement\", \"None\", TriageCategory.URGENT, \"Appendicitis\"),\n",
+                "    PatientScenario(\"URG_02\", \"55M Foot Infection\", {\"hr\": 92, \"bp_sys\": 145, \"bp_dia\": 88, \"rr\": 16, \"spo2\": 97, \"temp\": 37.4, \"avpu\": \"A\"}, \"Non-healing foot wound for 2 weeks, getting red\", \"T2DM\", TriageCategory.URGENT, \"Diabetic foot\"),\n",
+                "\n",
+                "    # STANDARD (4)\n",
+                "    PatientScenario(\"STD_01\", \"22M Ankle Sprain\", {\"hr\": 75, \"bp_sys\": 125, \"bp_dia\": 80, \"rr\": 14, \"spo2\": 99, \"temp\": 36.8, \"avpu\": \"A\"}, \"Twisted ankle playing football, swelling\", \"None\", TriageCategory.STANDARD, \"Sprain\"),\n",
+                "    \n",
+                "    # NON-URGENT (5)\n",
+                "    PatientScenario(\"NU_01\", \"24F Sore Throat\", {\"hr\": 78, \"bp_sys\": 118, \"bp_dia\": 72, \"rr\": 14, \"spo2\": 99, \"temp\": 37.8, \"avpu\": \"A\"}, \"Sore throat for 3 days, can swallow fine\", \"None\", TriageCategory.NON_URGENT, \"Minor illness\"),\n",
+                "    PatientScenario(\"NU_02\", \"60M Med Refill\", {\"hr\": 70, \"bp_sys\": 120, \"bp_dia\": 80, \"rr\": 12, \"spo2\": 99, \"temp\": 36.8, \"avpu\": \"A\"}, \"Needs insulin refill, lost bag on bus\", \"T2DM\", TriageCategory.NON_URGENT, \"Admin task\"),\n",
+                "]"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "inference_logic"
+            },
+            "outputs": [],
+            "source": [
+                "# EXACT Training Prompt Format\n",
+                "TRAINING_INSTRUCTION = \"You are an expert A&E Triage Nurse using the Manchester Triage System. Assess the following patient and provide your triage decision with clinical reasoning.\"\n",
+                "\n",
+                "def parse_demographics(desc):\n",
+                "    # Extract \"72M\" or \"65F\" from start of string\n",
+                "    match = re.search(r\"(\\d+)([MF])\", desc)\n",
+                "    if match:\n",
+                "        return int(match.group(1)), \"Male\" if match.group(2) == \"M\" else \"Female\"\n",
+                "    return \"Unknown\", \"Unknown\"\n",
+                "\n",
+                "def format_input(c):\n",
+                "    age, gender = parse_demographics(c.desc)\n",
+                "    \n",
+                "    # CRITICAL: Reconstruct the exact dictionary format used in training\n",
+                "    history_dict = {\n",
+                "        'age': age,\n",
+                "        'gender': gender,\n",
+                "        'relevant_PMH': c.history,\n",
+                "        'time_course': 'See complaint'\n",
+                "    }\n",
+                "    \n",
+                "    # Force consistent key order if needed (Python dicts preserve order now, but just in case)\n",
+                "    \n",
+                "    return f\"\"\"PATIENT PRESENTING TO A&E TRIAGE\n",
+                "\n",
+                "Chief Complaint: \"{c.complaint}\"\n",
+                "\n",
+                "Vitals:\n",
+                "- HR: {c.vitals.get('hr')} bpm\n",
+                "- BP: {c.vitals.get('bp_sys')}/{c.vitals.get('bp_dia')} mmHg\n",
+                "- SpO2: {c.vitals.get('spo2')}%\n",
+                "- RR: {c.vitals.get('rr')} /min\n",
+                "- Temp: {c.vitals.get('temp')}C\n",
+                "- AVPU: {c.vitals.get('avpu')}\n",
+                "\n",
+                "History: {history_dict}\n",
+                "\n",
+                "WAITING ROOM: 12 patients | AVAILABLE BEDS: 4\n",
+                "\n",
+                "What is your triage decision?\"\"\"\n",
+                "\n",
+                "def generate_prediction(s):\n",
+                "    alpaca_prompt = f\"\"\"Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n",
+                "\n",
+                "### Instruction:\n",
+                "{TRAINING_INSTRUCTION}\n",
+                "\n",
+                "### Input:\n",
+                "{format_input(s)}\n",
+                "\n",
+                "### Response:\n",
+                "\"\"\"\n",
+                "    \n",
+                "    inputs = tokenizer(\n",
+                "        [alpaca_prompt],\n",
+                "        return_tensors=\"pt\",\n",
+                "    ).to(\"cuda\")\n",
+                "    \n",
+                "    outputs = model.generate(**inputs, max_new_tokens=256, use_cache=True)\n",
+                "    full_text = tokenizer.batch_decode(outputs)[0]\n",
+                "    response = full_text.split(\"### Response:\")[-1].replace(\"<|eot_id|>\", \"\").strip()\n",
+                "    return response"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "section_1_quantitative"
+            },
+            "outputs": [],
+            "source": [
+                "def extract_category(text):\n",
+                "    try:\n",
+                "        # Prioritize explicit \"Category: X\"\n",
+                "        match = re.search(r\"Category:\\s*(\\d)\", text)\n",
+                "        if match: return int(match.group(1))\n",
+                "        # Fallback\n",
+                "        match = re.search(r'\\b([1-5])\\b', text)\n",
+                "        return int(match.group(1)) if match else -1\n",
+                "    except: return -1\n",
+                "\n",
+                "print(\"📊 SECTION 1: Quantitative Benchmark (15 Cases)\\n\")\n",
+                "results = []\n",
+                "\n",
+                "for s in tqdm(TEST_SCENARIOS, desc=\"Benchmarking\"):\n",
+                "    response = generate_prediction(s)\n",
+                "    pred_cat = extract_category(response)\n",
+                "    exp_cat = s.expected.value\n",
+                "    \n",
+                "    results.append({\n",
+                "        'id': s.id,\n",
+                "        'desc': s.desc,\n",
+                "        'exp': exp_cat,\n",
+                "        'pred': pred_cat,\n",
+                "        'match': pred_cat == exp_cat,\n",
+                "        'full_resp': response\n",
+                "    })\n",
+                "    \n",
+                "    icon = \"✅\" if pred_cat == exp_cat else \"❌\"\n",
+                "    print(f\"{icon} {s.id}: Pred={pred_cat} | Exp={exp_cat}\")\n",
+                "\n",
+                "df_quant = pd.DataFrame(results)\n",
+                "acc = df_quant['match'].mean() * 100\n",
+                "print(f\"\\n🏆 QUANTITATIVE ACCURACY: {acc:.1f}%\")"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {
+                "id": "section_2_judge"
+            },
+            "outputs": [],
+            "source": [
+                "print(\"\\n⚖️ SECTION 2: LLM-as-Judge Evaluation (Qualitative)\\n\")\n",
+                "\n",
+                "JUDGE_SYS_PROMPT = \"\"\"You are a Senior Clinical Auditor. Evaluate the AI Nurse's response.\n",
+                "Criteria:\n",
+                "1. Accuracy (1-5): Correct triage category?\n",
+                "2. Reasoning (1-5): Sound clinical logic?\n",
+                "3. Safety (PASS/FAIL): Did it miss a life-threat?\n",
+                "Return JSON: {\"accuracy\": int, \"reasoning\": int, \"safety\": \"str\", \"critique\": \"str\"}\"\"\"\n",
+                "\n",
+                "judge_results = []\n",
+                "\n",
+                "for row in results:\n",
+                "    user_prompt = f\"SCENARIO: {row['desc']} (Exp: Cat {row['exp']})\\nAI RESPONSE:\\n{row['full_resp']}\"\n",
+                "    \n",
+                "    eval_res = {}\n",
+                "    try:\n",
+                "        # Try GPT-5.2 first\n",
+                "        client = openai.OpenAI()\n",
+                "        completion = client.chat.completions.create(\n",
+                "            model=\"gpt-5.2\",\n",
+                "            messages=[{\"role\": \"system\", \"content\": JUDGE_SYS_PROMPT}, {\"role\": \"user\", \"content\": user_prompt}],\n",
+                "            response_format={\"type\": \"json_object\"}, \n",
+                "            temperature=0\n",
+                "        )\n",
+                "        eval_res = json.loads(completion.choices[0].message.content)\n",
+                "    except Exception as e:\n",
+                "        # Fallback to Gemini\n",
+                "        try:\n",
+                "            model_gem = genai.GenerativeModel('gemini-3.0-pro-preview', system_instruction=JUDGE_SYS_PROMPT)\n",
+                "            resp = model_gem.generate_content(user_prompt)\n",
+                "            text = resp.text.split(\"```json\")[-1].split(\"```\")[0] if \"```\" in resp.text else resp.text\n",
+                "            eval_res = json.loads(text)\n",
+                "        except:\n",
+                "             eval_res = {\"accuracy\": 0, \"safety\": \"ERROR\", \"critique\": \"Judge Failed\"}\n",
+                "\n",
+                "    print(f\"📝 {row['id']}: Safety={eval_res.get('safety')} | Judge Acc={eval_res.get('accuracy')}/5\")\n",
+                "    print(f\"   Critique: {eval_res.get('critique')[:100]}...\")\n",
+                "    \n",
+                "    row.update(eval_res)\n",
+                "    judge_results.append(row)\n",
+                "\n",
+                "df_final = pd.DataFrame(judge_results)\n",
+                "df_final.to_csv(\"unified_validation_results.csv\", index=False)\n",
+                "print(\"\\n✅ All Results saved to unified_validation_results.csv\")"
+            ]
+        }
+    ],
+    "metadata": {
+        "accelerator": "GPU",
+        "colab": {
+            "gpuType": "A100",
+            "provenance": []
+        },
+        "kernelspec": {
+            "display_name": "Python 3",
+            "name": "python3"
+        },
+        "language_info": {
+            "name": "python"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 0
+}

nursesim_rl/__init__.py

@@ -5,6 +5,12 @@ OpenEnv Challenge Entry - 2026
 
 from .triage_env import TriageEnv
 from .patient_generator import PatientGenerator
+from .semantic_wrapper import NurseEmbedWrapper, make_semantic_triage_env
 
-__version__ = "0.1.0"
-__all__ = ["TriageEnv", "PatientGenerator"]
+__version__ = "0.2.0"
+__all__ = [
+    "TriageEnv", 
+    "PatientGenerator",
+    "NurseEmbedWrapper",
+    "make_semantic_triage_env",
+]

nursesim_rl/semantic_wrapper.py

@@ -0,0 +1,169 @@
+"""
+NurseEmbedWrapper: A Gymnasium wrapper that converts text observations to NurseEmbed vectors.
+
+This enables Language-Conditioned Reinforcement Learning for nursing scenarios.
+"""
+
+import gymnasium as gym
+from gymnasium import spaces
+import numpy as np
+from typing import Any, Dict, Tuple
+
+# Lazy load NurseEmbed to avoid import errors if not available
+_embed_model = None
+
+def _get_embed_model():
+    """Lazy load the NurseEmbed model."""
+    global _embed_model
+    if _embed_model is None:
+        try:
+            from sentence_transformers import SentenceTransformer
+            _embed_model = SentenceTransformer("NurseCitizenDeveloper/NurseEmbed-300M")
+            print("[OK] NurseEmbed model loaded successfully")
+        except Exception as e:
+            print(f"[WARN] NurseEmbed not available: {e}")
+            # Fallback to random embeddings for testing
+            _embed_model = "fallback"
+    return _embed_model
+
+
+class NurseEmbedWrapper(gym.Wrapper):
+    """
+    Wraps a TriageEnv and converts text observations to NurseEmbed vectors.
+    Also flattens the Dict action space to MultiDiscrete for SB3 compatibility.
+    
+    Instead of the agent seeing:
+        {"chief_complaint": "Chest pain radiating to arm...", "vitals": {...}, ...}
+    
+    It sees:
+        np.array([...390 dimensional semantic vector...])
+    
+    The vector encodes the MEANING of the clinical presentation.
+    """
+    
+    EMBED_DIM = 384  # NurseEmbed-300M outputs 384D vectors (nomic base)
+    
+    def __init__(self, env: gym.Env, use_vitals: bool = True):
+        """
+        Args:
+            env: The underlying TriageEnv
+            use_vitals: Whether to append vitals to the embedding
+        """
+        super().__init__(env)
+        
+        self.use_vitals = use_vitals
+        self.model = _get_embed_model()
+        
+        # Calculate observation dimension
+        obs_dim = self.EMBED_DIM
+        if use_vitals:
+            obs_dim += 6  # HR, BP_sys, BP_dia, SpO2, RR, Temp
+        
+        # Override observation space to be a flat Box
+        self.observation_space = spaces.Box(
+            low=-np.inf,
+            high=np.inf,
+            shape=(obs_dim,),
+            dtype=np.float32
+        )
+        
+        # Flatten action space for SB3 compatibility
+        # Original: Dict({'triage_category': Discrete(5, start=1), 'intervention': Discrete(7)})
+        # New: MultiDiscrete([5, 7]) where first dim is category (0-4, add 1 later) and second is intervention
+        self.action_space = spaces.MultiDiscrete([5, 7])
+        
+        # Cache for embeddings (same text -> same embedding)
+        self._embedding_cache: Dict[str, np.ndarray] = {}
+    
+    def reset(self, **kwargs) -> Tuple[np.ndarray, Dict]:
+        """Reset and convert observation."""
+        obs, info = self.env.reset(**kwargs)
+        return self._convert_observation(obs), info
+    
+    def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, bool, Dict]:
+        """Convert action, step, convert observation."""
+        # Convert flat action [category_idx, intervention_idx] to Dict
+        dict_action = {
+            "triage_category": int(action[0]) + 1,  # 0-4 -> 1-5
+            "intervention": int(action[1])
+        }
+        obs, reward, terminated, truncated, info = self.env.step(dict_action)
+        return self._convert_observation(obs), reward, terminated, truncated, info
+    
+    def _convert_observation(self, obs: Dict) -> np.ndarray:
+        """Convert Dict observation to semantic vector."""
+        # Build text representation
+        text = self._build_clinical_text(obs)
+        
+        # Get embedding (with caching)
+        embedding = self._get_embedding(text)
+        
+        # Optionally append vitals
+        if self.use_vitals:
+            vitals_vector = self._extract_vitals(obs)
+            embedding = np.concatenate([embedding, vitals_vector])
+        
+        return embedding.astype(np.float32)
+    
+    def _build_clinical_text(self, obs: Dict) -> str:
+        """Build a clinical description from the observation."""
+        complaint = obs.get("chief_complaint", "Unknown complaint")
+        history = obs.get("history", "")
+        
+        # Create a rich clinical description
+        text = f"Patient presents with: {complaint}. Clinical history: {history}."
+        
+        # Add vitals context as text for semantic understanding
+        vitals = obs.get("vitals", {})
+        if vitals:
+            text += f" Vital signs: HR {vitals.get('hr', 'N/A')}, "
+            text += f"BP {vitals.get('bp_sys', 'N/A')}/{vitals.get('bp_dia', 'N/A')}, "
+            text += f"SpO2 {vitals.get('spo2', 'N/A')}%, "
+            text += f"RR {vitals.get('rr', 'N/A')}, "
+            text += f"AVPU {vitals.get('avpu', 'A')}."
+        
+        return text
+    
+    def _get_embedding(self, text: str) -> np.ndarray:
+        """Get embedding with caching."""
+        if text in self._embedding_cache:
+            return self._embedding_cache[text]
+        
+        if self.model == "fallback":
+            # Fallback: deterministic pseudo-random embedding based on text hash
+            np.random.seed(hash(text) % 2**32)
+            embedding = np.random.randn(self.EMBED_DIM)
+        else:
+            embedding = self.model.encode(text, normalize_embeddings=True)
+        
+        self._embedding_cache[text] = embedding
+        return embedding
+    
+    def _extract_vitals(self, obs: Dict) -> np.ndarray:
+        """Extract vitals as a normalized vector."""
+        vitals = obs.get("vitals", {})
+        return np.array([
+            vitals.get("hr", 70) / 200.0,  # Normalize HR
+            vitals.get("bp_sys", 120) / 200.0,
+            vitals.get("bp_dia", 80) / 150.0,
+            vitals.get("spo2", 98) / 100.0,
+            vitals.get("rr", 16) / 40.0,
+            (vitals.get("temp", 37.0) - 35) / 5.0,  # Normalize temp
+        ], dtype=np.float32)
+
+
+def make_semantic_triage_env(seed: int = None, **kwargs) -> gym.Env:
+    """Factory function to create a semantically-aware triage environment."""
+    from nursesim_rl import TriageEnv
+    
+    base_env = TriageEnv(seed=seed, **kwargs)
+    wrapped_env = NurseEmbedWrapper(base_env, use_vitals=True)
+    
+    return wrapped_env
+
+
+# Register wrapped version
+gym.register(
+    id="NurseSim-SemanticTriage-v0",
+    entry_point="nursesim_rl.semantic_wrapper:make_semantic_triage_env",
+)

test_semantic.py

@@ -0,0 +1,64 @@
+"""
+Test script for the Semantic RL Environment.
+Validates that the NurseEmbedWrapper works correctly.
+"""
+
+import numpy as np
+from nursesim_rl import TriageEnv, NurseEmbedWrapper
+
+def test_semantic_wrapper():
+    print("=" * 60)
+    print("Testing NurseEmbed Semantic RL Wrapper")
+    print("=" * 60)
+    
+    # Create base environment
+    print("\n[1] Creating base TriageEnv...")
+    base_env = TriageEnv(max_steps=10, seed=42)
+    
+    # Wrap with NurseEmbed
+    print("[2] Wrapping with NurseEmbedWrapper...")
+    semantic_env = NurseEmbedWrapper(base_env, use_vitals=True)
+    
+    # Check observation space
+    print(f"\n[3] Observation Space Check:")
+    print(f"    Base Env: {type(base_env.observation_space)}")
+    print(f"    Semantic Env: {semantic_env.observation_space}")
+    print(f"    Expected shape: (390,) [384 embed + 6 vitals]")
+    
+    # Reset and get observation
+    print("\n[4] Resetting environment...")
+    obs, info = semantic_env.reset(seed=42)
+    
+    print(f"    Observation type: {type(obs)}")
+    print(f"    Observation shape: {obs.shape}")
+    print(f"    Observation range: [{obs.min():.3f}, {obs.max():.3f}]")
+    
+    # Take a step
+    print("\n[5] Taking a step with action (Cat=3, Intervention=2)...")
+    action = {"triage_category": 3, "intervention": 2}
+    obs2, reward, terminated, truncated, info = semantic_env.step(action)
+    
+    print(f"    Reward: {reward}")
+    print(f"    Terminated: {terminated}")
+    print(f"    New observation shape: {obs2.shape}")
+    
+    # Verify embedding is meaningful (not just zeros)
+    print("\n[6] Embedding quality check:")
+    embed_part = obs[:384]  # First 384 dims are the embedding
+    print(f"    Embedding L2 norm: {np.linalg.norm(embed_part):.3f}")
+    print(f"    Embedding is normalized: {abs(np.linalg.norm(embed_part) - 1.0) < 0.1}")
+    
+    # Test caching
+    print("\n[7] Testing embedding cache...")
+    obs3, _ = semantic_env.reset(seed=42)  # Same seed = same patient
+    cache_hit = np.allclose(obs[:384], obs3[:384])
+    print(f"    Cache working: {cache_hit}")
+    
+    print("\n" + "=" * 60)
+    print("ALL TESTS PASSED!")
+    print("=" * 60)
+    
+    return True
+
+if __name__ == "__main__":
+    test_semantic_wrapper()

train_semantic_agent.py

@@ -0,0 +1,143 @@
+"""
+Train a PPO Agent on the Semantic Triage Environment.
+
+This script trains an agent that learns from NurseEmbed-encoded clinical observations.
+"""
+
+import os
+import numpy as np
+from datetime import datetime
+
+# Stable Baselines 3
+from stable_baselines3 import PPO
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.callbacks import EvalCallback, BaseCallback
+from stable_baselines3.common.vec_env import DummyVecEnv
+
+# Our environment
+from nursesim_rl import TriageEnv, NurseEmbedWrapper
+
+
+class PrintProgressCallback(BaseCallback):
+    """Simple callback to print training progress."""
+    
+    def __init__(self, print_freq: int = 1000, verbose: int = 0):
+        super().__init__(verbose)
+        self.print_freq = print_freq
+        
+    def _on_step(self) -> bool:
+        if self.n_calls % self.print_freq == 0:
+            # Get recent episode rewards if available
+            if len(self.model.ep_info_buffer) > 0:
+                mean_reward = np.mean([ep['r'] for ep in self.model.ep_info_buffer])
+                mean_length = np.mean([ep['l'] for ep in self.model.ep_info_buffer])
+                print(f"Step {self.n_calls}: Mean Reward = {mean_reward:.2f}, Mean Length = {mean_length:.1f}")
+        return True
+
+
+def make_semantic_env():
+    """Factory function for creating the semantic environment."""
+    base_env = TriageEnv(max_steps=50, max_patients=20)
+    semantic_env = NurseEmbedWrapper(base_env, use_vitals=True)
+    return semantic_env
+
+
+def train_semantic_agent(
+    total_timesteps: int = 50000,
+    save_path: str = "models/semantic_ppo",
+    log_dir: str = "logs/semantic_ppo",
+):
+    """Train a PPO agent on the semantic triage environment."""
+    
+    print("=" * 60)
+    print("SEMANTIC RL TRAINING")
+    print(f"Started: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
+    print("=" * 60)
+    
+    # Create directories
+    os.makedirs(save_path, exist_ok=True)
+    os.makedirs(log_dir, exist_ok=True)
+    
+    # Create vectorized environment
+    print("\n[1] Creating Semantic Environment...")
+    env = DummyVecEnv([make_semantic_env])
+    
+    print(f"    Observation space: {env.observation_space}")
+    print(f"    Action space: {env.action_space}")
+    
+    # Create evaluation environment
+    eval_env = DummyVecEnv([make_semantic_env])
+    
+    # Create the PPO agent
+    print("\n[2] Initializing PPO Agent...")
+    model = PPO(
+        "MlpPolicy",
+        env,
+        learning_rate=3e-4,
+        n_steps=2048,
+        batch_size=64,
+        n_epochs=10,
+        gamma=0.99,
+        gae_lambda=0.95,
+        clip_range=0.2,
+        ent_coef=0.01,
+        verbose=0,
+        tensorboard_log=log_dir,
+    )
+    
+    print(f"    Policy architecture: {model.policy}")
+    
+    # Callbacks
+    progress_callback = PrintProgressCallback(print_freq=2000)
+    eval_callback = EvalCallback(
+        eval_env,
+        best_model_save_path=save_path,
+        log_path=log_dir,
+        eval_freq=5000,
+        n_eval_episodes=5,
+        deterministic=True,
+        verbose=0,
+    )
+    
+    # Train!
+    print(f"\n[3] Training for {total_timesteps:,} timesteps...")
+    print("-" * 60)
+    
+    model.learn(
+        total_timesteps=total_timesteps,
+        callback=[progress_callback, eval_callback],
+        progress_bar=True,
+    )
+    
+    # Save final model
+    final_path = os.path.join(save_path, "semantic_ppo_final")
+    model.save(final_path)
+    print(f"\n[4] Model saved to: {final_path}")
+    
+    # Quick evaluation
+    print("\n[5] Final Evaluation (10 episodes)...")
+    rewards = []
+    for i in range(10):
+        obs = eval_env.reset()
+        episode_reward = 0
+        done = False
+        while not done:
+            action, _ = model.predict(obs, deterministic=True)
+            obs, reward, done, info = eval_env.step(action)
+            episode_reward += reward[0]
+        rewards.append(episode_reward)
+    
+    print(f"    Mean Reward: {np.mean(rewards):.2f} +/- {np.std(rewards):.2f}")
+    print(f"    Best Episode: {np.max(rewards):.2f}")
+    print(f"    Worst Episode: {np.min(rewards):.2f}")
+    
+    print("\n" + "=" * 60)
+    print("TRAINING COMPLETE!")
+    print("=" * 60)
+    
+    return model
+
+
+if __name__ == "__main__":
+    # Run with reduced timesteps for quick demo
+    train_semantic_agent(total_timesteps=20000)

viz_semantic.py

@@ -0,0 +1,173 @@
+"""
+Visualize "What the Agent Sees" - Semantic Embedding Projection
+
+This script generates a 2D t-SNE visualization of how the RL agent
+perceives clinical observations through NurseEmbed semantic vectors.
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+from sklearn.decomposition import PCA
+import os
+
+from nursesim_rl import TriageEnv, NurseEmbedWrapper
+
+
+def collect_observations(n_observations: int = 200, seed: int = 42):
+    """Collect observations and their true categories."""
+    
+    print(f"[1] Collecting {n_observations} observations...")
+    
+    # Create semantic environment
+    base_env = TriageEnv(max_steps=100, max_patients=50, seed=seed)
+    env = NurseEmbedWrapper(base_env, use_vitals=True)
+    
+    observations = []
+    categories = []
+    complaints = []
+    
+    obs, info = env.reset(seed=seed)
+    
+    for i in range(n_observations):
+        # Store observation and metadata
+        observations.append(obs.copy())
+        
+        # Get true category from wrapped env
+        if base_env.current_patient:
+            categories.append(base_env.current_patient.true_category)
+            complaints.append(base_env.current_patient.chief_complaint[:50])
+        else:
+            categories.append(3)  # Default
+            complaints.append("No patient")
+        
+        # Take a random action to move to next patient
+        action = env.action_space.sample()
+        obs, _, terminated, truncated, _ = env.step(action)
+        
+        if terminated or truncated:
+            obs, _ = env.reset()
+    
+    print(f"    Collected {len(observations)} observations")
+    
+    return np.array(observations), np.array(categories), complaints
+
+
+def visualize_embeddings(observations, categories, complaints, output_path="viz"):
+    """Create t-SNE visualization of embeddings."""
+    
+    print("[2] Computing t-SNE projection...")
+    
+    # Extract just the embedding part (first 384 dims)
+    embeddings = observations[:, :384]
+    
+    # First reduce with PCA for speed (384 -> 50)
+    pca = PCA(n_components=50)
+    embeddings_pca = pca.fit_transform(embeddings)
+    print(f"    PCA explained variance: {sum(pca.explained_variance_ratio_):.2%}")
+    
+    # Then t-SNE to 2D
+    tsne = TSNE(n_components=2, perplexity=30, random_state=42, max_iter=1000)
+    embeddings_2d = tsne.fit_transform(embeddings_pca)
+    
+    print("[3] Creating visualization...")
+    
+    # Create figure
+    fig, ax = plt.subplots(figsize=(12, 10))
+    
+    # Color map for triage categories
+    colors = {
+        1: '#FF0000',  # Immediate - Red
+        2: '#FF8C00',  # Very Urgent - Orange
+        3: '#FFD700',  # Urgent - Yellow
+        4: '#32CD32',  # Standard - Green
+        5: '#1E90FF',  # Non-urgent - Blue
+    }
+    
+    category_names = {
+        1: 'P1: Immediate',
+        2: 'P2: Very Urgent',
+        3: 'P3: Urgent',
+        4: 'P4: Standard',
+        5: 'P5: Non-urgent'
+    }
+    
+    # Plot each category
+    for cat in sorted(set(categories)):
+        mask = categories == cat
+        ax.scatter(
+            embeddings_2d[mask, 0],
+            embeddings_2d[mask, 1],
+            c=colors.get(cat, '#888888'),
+            label=category_names.get(cat, f'Category {cat}'),
+            alpha=0.7,
+            s=100,
+            edgecolors='white',
+            linewidths=0.5
+        )
+    
+    # Styling
+    ax.set_title(
+        'What the Agent Sees: Semantic Patient Clusters\n'
+        '(t-SNE projection of NurseEmbed vectors)',
+        fontsize=14, fontweight='bold'
+    )
+    ax.set_xlabel('Semantic Dimension 1', fontsize=11)
+    ax.set_ylabel('Semantic Dimension 2', fontsize=11)
+    
+    # Legend
+    ax.legend(loc='upper right', title='Triage Category', fontsize=10)
+    
+    # Clean up axes
+    ax.spines['top'].set_visible(False)
+    ax.spines['right'].set_visible(False)
+    ax.grid(True, alpha=0.3)
+    
+    # Add annotation
+    ax.text(
+        0.02, 0.02,
+        'Patients with similar clinical presentations\ncluster together in semantic space.',
+        transform=ax.transAxes,
+        fontsize=9,
+        verticalalignment='bottom',
+        style='italic',
+        bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5)
+    )
+    
+    # Save
+    os.makedirs(output_path, exist_ok=True)
+    output_file = os.path.join(output_path, 'semantic_clusters.png')
+    plt.tight_layout()
+    plt.savefig(output_file, dpi=150, bbox_inches='tight')
+    print(f"[4] Saved to: {output_file}")
+    
+    # Also show statistics
+    print("\n[5] Cluster Statistics:")
+    for cat in sorted(set(categories)):
+        mask = categories == cat
+        center = embeddings_2d[mask].mean(axis=0)
+        print(f"    {category_names[cat]}: {mask.sum()} patients, center at ({center[0]:.1f}, {center[1]:.1f})")
+    
+    return output_file
+
+
+def main():
+    print("=" * 60)
+    print("SEMANTIC EMBEDDING VISUALIZATION")
+    print("'What the Agent Sees'")
+    print("=" * 60 + "\n")
+    
+    # Collect data
+    observations, categories, complaints = collect_observations(n_observations=150)
+    
+    # Visualize
+    output_file = visualize_embeddings(observations, categories, complaints)
+    
+    print("\n" + "=" * 60)
+    print("VISUALIZATION COMPLETE!")
+    print(f"Open: {os.path.abspath(output_file)}")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()