Add local Nemotron validation UI

gradio_pharmacopilot_demo.py

@@ -51,11 +51,15 @@ INVENTORY_PATH = data_path("inventory.json")
 
 MODEL_ID = os.getenv("PHARMACOPILOT_MODEL_ID", "openbmb/MiniCPM-V-4_5")
 LIVE_GPU_OCR = os.getenv("PHARMACOPILOT_LIVE_GPU_OCR", "1").lower() not in {"0", "false", "no"}
+LIVE_NEMOTRON = os.getenv("PHARMACOPILOT_LIVE_NEMOTRON", "1").lower() not in {"0", "false", "no"}
+NEMOTRON_MODEL_ID = os.getenv("NEMOTRON_MODEL_ID", "nvidia/NVIDIA-Nemotron-Nano-9B-v2")
 DEMO_OCR_TEXT = "Neuoxen"
 DEMO_PROMPT = "Read the handwritten medicine name in the image. Return only the text."
 ACCEPTANCE_THRESHOLD = int(os.getenv("PHARMACOPILOT_ACCEPTANCE_THRESHOLD", "75"))
 OCR_MODEL = None
 OCR_TOKENIZER = None
+NEMOTRON_MODEL = None
+NEMOTRON_TOKENIZER = None
 
 
 def load_json(path: Path, fallback: Any) -> Any:
@@ -213,6 +217,183 @@ def first_strength(strength: str) -> str:
     return strength.split(",")[0].strip()
 
 
+def fallback_prescription_plan(
+    ocr_text: str,
+    medicine: dict[str, Any],
+    display_name: str,
+    confidence: int,
+    note: str = "Nemotron did not run",
+) -> dict[str, Any]:
+    accepted = confidence >= ACCEPTANCE_THRESHOLD
+    return {
+        "status": "needs_review" if not accepted else "retrieval_only",
+        "medicine_name": display_name if accepted else "Needs review",
+        "canonical_name": medicine.get("name", "Unknown") if accepted else f"Suggestion: {medicine.get('name', 'Unknown')}",
+        "dose": first_strength(medicine.get("strength", "")) if accepted else "Not confirmed",
+        "route": "Not specified",
+        "timing": "Not specified",
+        "frequency": "Not specified",
+        "duration": "Not specified",
+        "instructions": "Pharmacist review required before dispensing.",
+        "validation_note": note,
+        "ocr_text": ocr_text,
+    }
+
+
+def release_model_memory() -> None:
+    try:
+        import gc
+        import torch
+
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    except Exception:
+        return
+
+
+def unload_ocr_model() -> None:
+    global OCR_MODEL, OCR_TOKENIZER
+    OCR_MODEL = None
+    OCR_TOKENIZER = None
+    release_model_memory()
+
+
+def unload_nemotron_model() -> None:
+    global NEMOTRON_MODEL, NEMOTRON_TOKENIZER
+    NEMOTRON_MODEL = None
+    NEMOTRON_TOKENIZER = None
+    release_model_memory()
+
+
+def extract_json_object(text: str) -> dict[str, Any]:
+    cleaned = str(text or "").strip()
+    cleaned = cleaned.removeprefix("```json").removeprefix("```").removesuffix("```").strip()
+    start = cleaned.find("{")
+    end = cleaned.rfind("}")
+    if start >= 0 and end > start:
+        cleaned = cleaned[start : end + 1]
+    return json.loads(cleaned)
+
+
+def validate_with_nemotron(
+    ocr_text: str,
+    medicine: dict[str, Any],
+    display_name: str,
+    confidence: int,
+    retrieval_candidates: list[dict[str, Any]],
+) -> dict[str, Any]:
+    global NEMOTRON_MODEL, NEMOTRON_TOKENIZER
+
+    if not LIVE_NEMOTRON:
+        return fallback_prescription_plan(
+            ocr_text, medicine, display_name, confidence, "Local Nemotron validation is disabled"
+        )
+
+    validation_payload = {
+        "ocr_text": ocr_text,
+        "retrieved_display_name": display_name,
+        "retrieved_canonical_name": medicine.get("name", "Unknown"),
+        "retrieval_confidence": confidence,
+        "strength": first_strength(medicine.get("strength", "")),
+        "category": medicine.get("category", "Unknown"),
+        "top_candidates": [
+            {
+                "display_name": item["label"],
+                "canonical_name": item["medicine"]["name"],
+                "score": round(item["score"] * 100),
+            }
+            for item in retrieval_candidates[:3]
+        ],
+    }
+    prompt = f"""
+You are a pharmacy prescription validation assistant.
+
+Input JSON:
+{json.dumps(validation_payload, ensure_ascii=False)}
+
+Task:
+1. Decide whether the retrieved medicine is safe to accept.
+2. Translate the prescription into a clean pharmacy instruction row.
+3. Do not invent dose/timing/duration if it is not visible or inferable.
+
+Return ONLY valid JSON with these keys:
+status: one of validated, needs_review
+medicine_name
+canonical_name
+dose
+route
+timing
+frequency
+duration
+instructions
+validation_note
+ocr_text
+"""
+    try:
+        import torch
+        from transformers import AutoModelForCausalLM, AutoTokenizer
+
+        if NEMOTRON_MODEL is None or NEMOTRON_TOKENIZER is None:
+            NEMOTRON_TOKENIZER = AutoTokenizer.from_pretrained(NEMOTRON_MODEL_ID, trust_remote_code=True)
+            NEMOTRON_MODEL = AutoModelForCausalLM.from_pretrained(
+                NEMOTRON_MODEL_ID,
+                trust_remote_code=True,
+                torch_dtype=torch.bfloat16,
+                device_map="auto",
+            ).eval()
+
+        messages = [{"role": "user", "content": prompt}]
+        if hasattr(NEMOTRON_TOKENIZER, "apply_chat_template"):
+            input_ids = NEMOTRON_TOKENIZER.apply_chat_template(
+                messages,
+                add_generation_prompt=True,
+                return_tensors="pt",
+            )
+        else:
+            input_ids = NEMOTRON_TOKENIZER(prompt, return_tensors="pt").input_ids
+
+        device = next(NEMOTRON_MODEL.parameters()).device
+        input_ids = input_ids.to(device)
+        with torch.inference_mode():
+            output_ids = NEMOTRON_MODEL.generate(
+                input_ids,
+                do_sample=False,
+                temperature=0.0,
+                top_p=1.0,
+                max_new_tokens=320,
+                pad_token_id=NEMOTRON_TOKENIZER.eos_token_id,
+            )
+        generated = output_ids[0][input_ids.shape[-1] :]
+        content = NEMOTRON_TOKENIZER.decode(generated, skip_special_tokens=True).strip()
+        plan = extract_json_object(content)
+        if plan.get("status") not in {"validated", "needs_review"}:
+            plan["status"] = "needs_review"
+        if confidence < ACCEPTANCE_THRESHOLD:
+            plan["status"] = "needs_review"
+            plan["validation_note"] = (
+                f"Retrieval confidence {confidence}% is below the {ACCEPTANCE_THRESHOLD}% acceptance threshold"
+            )
+        return {
+            **fallback_prescription_plan(
+                ocr_text,
+                medicine,
+                display_name,
+                confidence,
+                f"Validated by local {NEMOTRON_MODEL_ID}",
+            ),
+            **plan,
+        }
+    except Exception as exc:
+        return fallback_prescription_plan(
+            ocr_text,
+            medicine,
+            display_name,
+            confidence,
+            f"Local Nemotron failed: {exc}",
+        )
+
+
 def load_kpi_metrics(searches: int = 0) -> str:
     metrics_path = ROOT / "training" / "baseline_eval" / "minicpm_v_4_5" / "baseline_minicpm_v_4_5_metrics.json"
     fallback_path = ROOT / "training" / "baseline_eval" / "minicpm_v_4_5" / "baseline_minicpm_v_4_5_report.md"
@@ -273,13 +454,19 @@ def confidence_gauge(confidence: int = 97):
     return fig
 
 
-def pipeline_html(stage: int = 0) -> str:
+def pipeline_html(stage: int = 0, validation_status: str = "waiting") -> str:
+    validation_label = {
+        "validated": "Nemotron Validated",
+        "needs_review": "Needs Review",
+        "retrieval_only": "Retrieval Only",
+        "waiting": "Awaiting Analysis",
+    }.get(validation_status, "Nemotron Review")
     steps = [
-        ("Prescription", "captured"),
-        ("MiniCPM OCR", "completed"),
-        ("Retrieval Engine", "searched"),
-        ("Validation", "checked"),
-        ("Medicine Found", "ready"),
+        ("Prescription", "uploaded"),
+        ("MiniCPM OCR", "ran on image"),
+        ("Retrieval Engine", "ranked candidates"),
+        (validation_label, "returned a decision"),
+        ("Pharmacy View", "prepared"),
     ]
     cards = []
     logs = []
@@ -297,7 +484,7 @@ def pipeline_html(stage: int = 0) -> str:
             logs.append(f"<li>✓ {title} {status}</li>")
     return f"""
     <div class="pipeline">
-      <div class="pipeline-title">Live AI Pipeline</div>
+      <div class="pipeline-title">Actual Run Trace</div>
       <div class="flow">{''.join(cards)}</div>
       <ul class="logs">{''.join(logs)}</ul>
     </div>
@@ -310,15 +497,16 @@ def medicine_details_html(
     ocr_text: str,
     display_name: str,
     confidence: int,
+    plan: dict[str, Any],
 ) -> str:
-    accepted = confidence >= ACCEPTANCE_THRESHOLD
+    accepted = plan.get("status") == "validated" and confidence >= ACCEPTANCE_THRESHOLD
     medicine_label = display_name if accepted else "Needs pharmacist review"
     generic_label = medicine.get("name", "Unknown") if accepted else f"Suggestion: {medicine.get('name', 'Unknown')}"
     strength_label = first_strength(medicine.get("strength", "")) if accepted else "Not confirmed"
     manufacturer_label = (medicine.get("manufacturer") or "Not listed") if accepted else "Not confirmed"
     category_label = medicine.get("category", "General") if accepted else "Not confirmed"
     price_label = "PKR 145" if accepted else "Not confirmed"
-    validation_label = "Pending review" if not accepted else "Accepted by retrieval threshold"
+    validation_label = plan.get("validation_note") or plan.get("status", "Not available")
     inventory_label = (
         f"Shelf {inventory['shelf']}, row {inventory['row']}"
         if accepted
@@ -347,6 +535,31 @@ def medicine_details_html(
     """
 
 
+def translated_prescription_html(plan: dict[str, Any]) -> str:
+    rows = [
+        ("Medicine", plan.get("medicine_name") or "Not confirmed"),
+        ("Canonical", plan.get("canonical_name") or "Not confirmed"),
+        ("Dose", plan.get("dose") or "Not specified"),
+        ("Route", plan.get("route") or "Not specified"),
+        ("When to take", plan.get("timing") or "Not specified"),
+        ("Pill timing", plan.get("frequency") or "Not specified"),
+        ("Duration", plan.get("duration") or "Not specified"),
+        ("Instructions", plan.get("instructions") or "Pharmacist review required"),
+    ]
+    row_html = "".join(f"<dt>{label}</dt><dd>{value}</dd>" for label, value in rows)
+    status = plan.get("status", "needs_review").replace("_", " ").title()
+    return f"""
+    <div class="translated-card">
+      <div class="translated-head">
+        <h3>Translated Prescription</h3>
+        <span class="status-pill">{status}</span>
+      </div>
+      <dl class="details translated-details">{row_html}</dl>
+      <p class="fine-print">Generated from OCR text and retrieval candidates. Confirm before dispensing.</p>
+    </div>
+    """
+
+
 def package_status_html(inventory: dict[str, Any], accepted: bool = True) -> str:
     if not accepted:
         return """
@@ -382,13 +595,19 @@ def candidates_html(candidates: list[dict[str, Any]]) -> str:
     """
 
 
-def ocr_compare_html(medicine: dict[str, Any], ocr_text: str, display_name: str, confidence: int) -> str:
-    corrected = display_name if confidence >= ACCEPTANCE_THRESHOLD else f"Needs review: {display_name}"
+def ocr_compare_html(
+    medicine: dict[str, Any],
+    ocr_text: str,
+    display_name: str,
+    confidence: int,
+    plan: dict[str, Any],
+) -> str:
+    corrected = display_name if plan.get("status") == "validated" else f"Needs review: {display_name}"
     return f"""
     <div class="compare-grid">
       <div><span>OCR Output</span><strong>{ocr_text}</strong></div>
       <div><span>AI Corrected</span><strong>{corrected}</strong></div>
-      <div><span>Canonical</span><strong>{medicine['name']}</strong></div>
+      <div><span>Canonical</span><strong>{medicine['name'] if plan.get('status') == 'validated' else 'Not confirmed'}</strong></div>
     </div>
     """
 
@@ -439,7 +658,7 @@ def run_minicpm_ocr(pil_image: Image.Image) -> str:
     return clean_prediction(raw_prediction) or raw_prediction.strip()
 
 
-@spaces.GPU(duration=60)
+@spaces.GPU(duration=300)
 def analyze_prescription(image, progress=gr.Progress()):
     global SESSION_SEARCHES
     if image is None:
@@ -453,17 +672,20 @@ def analyze_prescription(image, progress=gr.Progress()):
         time.sleep(0.15)
 
     ocr_text = run_minicpm_ocr(image)
+    unload_ocr_model()
 
     for pct, label in [
         (0.70, "Retrieval search over medicine aliases"),
-        (0.88, "Validation check"),
+        (0.88, "Nemotron prescription validation"),
         (1.00, "Result prepared"),
     ]:
         progress(pct, desc=label)
         time.sleep(0.25)
 
     medicine, candidates, display_name, confidence = find_medicine_from_ocr(ocr_text)
-    accepted = confidence >= ACCEPTANCE_THRESHOLD
+    plan = validate_with_nemotron(ocr_text, medicine, display_name, confidence, candidates)
+    unload_nemotron_model()
+    accepted = plan.get("status") == "validated" and confidence >= ACCEPTANCE_THRESHOLD
     inventory = get_inventory(medicine)
     image_path = resolve_asset_path(medicine.get("image_path"))
     package_image = str(image_path) if image_path and accepted else None
@@ -480,13 +702,14 @@ def analyze_prescription(image, progress=gr.Progress()):
 
     return (
         load_kpi_metrics(SESSION_SEARCHES),
-        pipeline_html(5),
-        medicine_details_html(medicine, inventory, ocr_text, display_name, confidence),
+        pipeline_html(5, plan.get("status", "needs_review")),
+        medicine_details_html(medicine, inventory, ocr_text, display_name, confidence, plan),
         package_image,
         package_status_html(inventory, accepted),
         confidence_gauge(confidence),
         candidates_html(candidates),
-        ocr_compare_html(medicine, ocr_text, display_name, confidence),
+        ocr_compare_html(medicine, ocr_text, display_name, confidence, plan),
+        translated_prescription_html(plan),
         gr.update(visible=True),
         gr.update(visible=True, interactive=accepted),
         state,
@@ -704,6 +927,40 @@ CSS = """
   background: #ffffff;
 }
 .compare-grid strong { color: var(--ink) !important; display: block; margin-top: 6px; font-size: 18px; }
+.translated-card {
+  border: 1px solid var(--line);
+  background: #ffffff;
+  border-radius: 8px;
+  padding: 18px;
+  margin-top: 12px;
+}
+.translated-head {
+  display: flex;
+  justify-content: space-between;
+  gap: 12px;
+  align-items: center;
+  margin-bottom: 12px;
+}
+.translated-head h3 { color: var(--ink) !important; margin: 0; font-size: 20px; }
+.status-pill {
+  background: #ebfbf5;
+  border: 1px solid #85d7bd;
+  color: #075f45;
+  border-radius: 999px;
+  padding: 5px 10px;
+  font-size: 12px;
+  font-weight: 800;
+}
+.translated-details {
+  grid-template-columns: 130px 1fr;
+}
+.fine-print {
+  border-top: 1px solid var(--line);
+  color: var(--muted) !important;
+  margin: 14px 0 0;
+  padding-top: 12px;
+  font-size: 13px;
+}
 .compact { margin-top: 0; }
 .gradio-container button.primary,
 .gradio-container button[variant="primary"] {
@@ -721,6 +978,7 @@ CSS = """
   .powered { text-align: left; margin-top: 10px; }
   .metric-row, .flow, .stock-card, .compare-grid { grid-template-columns: 1fr; }
   .details { grid-template-columns: 1fr; }
+  .translated-head { align-items: flex-start; flex-direction: column; }
 }
 """
 
@@ -776,6 +1034,7 @@ with gr.Blocks(title="PharmaCopilot") as demo:
             candidates = gr.HTML()
         gr.Markdown("### OCR vs Corrected")
         comparison = gr.HTML()
+        translated_prescription = gr.HTML()
         locate_btn = gr.Button("Locate Medicine", variant="primary", size="lg")
         locate_status = gr.Markdown()
 
@@ -806,6 +1065,7 @@ with gr.Blocks(title="PharmaCopilot") as demo:
             gauge,
             candidates,
             comparison,
+            translated_prescription,
             result_section,
             locate_btn,
             state,