Fallback Nemotron validation to NVIDIA NIM

gradio_pharmacopilot_demo.py

@@ -53,6 +53,9 @@ 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")
+NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY", "")
+NVIDIA_BASE_URL = os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1")
+NVIDIA_NIM_MODEL = os.getenv("NVIDIA_NIM_MODEL", "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"))
@@ -276,20 +279,13 @@ def extract_json_object(text: str) -> dict[str, Any]:
     return json.loads(cleaned)
 
 
-def validate_with_nemotron(
+def build_validation_prompt(
     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"
-        )
-
+) -> str:
     validation_payload = {
         "ocr_text": ocr_text,
         "retrieved_display_name": display_name,
@@ -306,7 +302,7 @@ def validate_with_nemotron(
             for item in retrieval_candidates[:3]
         ],
     }
-    prompt = f"""
+    return f"""
 You are a pharmacy prescription validation assistant.
 
 Input JSON:
@@ -316,6 +312,7 @@ 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.
+4. If OCR and retrieved medicine clearly disagree, return needs_review.
 
 Return ONLY valid JSON with these keys:
 status: one of validated, needs_review
@@ -330,6 +327,78 @@ instructions
 validation_note
 ocr_text
 """
+
+
+def validate_with_nvidia_nim(
+    prompt: str,
+    ocr_text: str,
+    medicine: dict[str, Any],
+    display_name: str,
+    confidence: int,
+) -> dict[str, Any]:
+    if not NVIDIA_API_KEY:
+        return fallback_prescription_plan(
+            ocr_text,
+            medicine,
+            display_name,
+            confidence,
+            "NVIDIA_API_KEY is not configured in the Space secrets",
+        )
+    try:
+        from openai import OpenAI
+
+        client = OpenAI(base_url=NVIDIA_BASE_URL, api_key=NVIDIA_API_KEY)
+        response = client.chat.completions.create(
+            model=NVIDIA_NIM_MODEL,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0,
+            top_p=1,
+            max_tokens=320,
+        )
+        content = response.choices[0].message.content or ""
+        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 NVIDIA NIM {NVIDIA_NIM_MODEL}",
+            ),
+            **plan,
+        }
+    except Exception as exc:
+        return fallback_prescription_plan(
+            ocr_text,
+            medicine,
+            display_name,
+            confidence,
+            f"NVIDIA NIM validation failed: {exc}",
+        )
+
+
+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"
+        )
+
+    prompt = build_validation_prompt(ocr_text, medicine, display_name, confidence, retrieval_candidates)
     try:
         import torch
         from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -385,13 +454,11 @@ ocr_text
             **plan,
         }
     except Exception as exc:
-        return fallback_prescription_plan(
-            ocr_text,
-            medicine,
-            display_name,
-            confidence,
-            f"Local Nemotron failed: {exc}",
-        )
+        nim_plan = validate_with_nvidia_nim(prompt, ocr_text, medicine, display_name, confidence)
+        if NVIDIA_API_KEY:
+            return nim_plan
+        nim_plan["validation_note"] = f"Local Nemotron failed: {exc}. NVIDIA_API_KEY is not configured."
+        return nim_plan
 
 
 def load_kpi_metrics(searches: int = 0) -> str:

requirements.txt

@@ -10,5 +10,4 @@ sentencepiece
 protobuf
 einops
 timm
-mamba-ssm>=2.2.5
-causal-conv1d>=1.5.0
+openai>=1.88.0