Update app.py

app.py

@@ -1,9 +1,8 @@
-# app.py
+```python
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
 import torch
-import re
 
 MODEL_ID = "Muhammadidrees/MedicalInsights"
 
@@ -12,151 +11,93 @@ MODEL_ID = "Muhammadidrees/MedicalInsights"
 # -----------------------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Try a few loading strategies so this works on GPU or CPU Spaces
 try:
-    # Preferred: let HF decide device placement (works for GPU-enabled Spaces)
     model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
 except Exception:
-    # Fallback: force CPU (slower but safe)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True
+    )
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1,
+)
 
-# Create pipeline
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+# -----------------------
+# Lookup Table for Biomarkers
+# -----------------------
+REFERENCE_RANGES = {
+    "Albumin": {"low": 3.5, "high": 5.5, "unit": "g/dL"},
+    "Creatinine_Male": {"low": 0.7, "high": 1.3, "unit": "mg/dL"},
+    "Creatinine_Female": {"low": 0.6, "high": 1.1, "unit": "mg/dL"},
+    "Glucose": {"low": 70, "high": 100, "unit": "mg/dL"},
+    "CRP": {"low": 0.3, "high": 10, "unit": "mg/L"},
+    "MCV": {"low": 80, "high": 100, "unit": "fL"},
+    "RDW": {"low": 11, "high": 15, "unit": "%"},
+    "WBC": {"low": 4000, "high": 11000, "unit": "/µL"},
+    "Lymphocytes": {"low": 20, "high": 40, "unit": "%"},
+    "ALP": {"low": 44, "high": 147, "unit": "U/L"},
+}
+
+def classify(value, biomarker, gender="Male"):
+    """Classify biomarker as Low, Normal, or High using lookup table."""
+    if biomarker == "Creatinine":
+        ref = REFERENCE_RANGES[f"Creatinine_{gender}"]
+    else:
+        ref = REFERENCE_RANGES[biomarker]
+
+    if value < ref["low"]:
+        return "Low"
+    elif value > ref["high"]:
+        return "High"
+    else:
+        return "Normal"
 
 # -----------------------
-# Helper: robust section splitter
+# Splitter helper (unchanged)
 # -----------------------
 def split_report(text):
-    """
-    Split model output into left (sections 1-4) and right (sections 5-6).
-    Accepts various markers for robustness.
-    """
-    # Normalize whitespace
     text = text.strip()
-    # Common markers that indicate tabular/insights section
-    markers = [
-        "5. Tabular Mapping",
-        "5. Tabular",
-        "Tabular Mapping",
-        "Tabular & AI Insights",
-        "📊 Tabular",
-        "## 5",
-    ]
-    # Find earliest marker occurrence
+    markers = ["5. Tabular Mapping", "Tabular Mapping", "📊 Tabular", "## 5"]
     idx = None
     for m in markers:
         pos = text.find(m)
-        if pos != -1:
-            if idx is None or pos < idx:
-                idx = pos
+        if pos != -1 and (idx is None or pos < idx):
+            idx = pos
     if idx is None:
-        # fallback: try splitting at "Enhanced AI Insights" or "Enhanced AI"
-        fallback = text.find("Enhanced AI Insights")
-        if fallback == -1:
-            fallback = text.find("Enhanced AI")
-        idx = fallback if fallback != -1 else None
-
-    if idx is None:
-        # couldn't find a split marker -> put everything in left
         return text, ""
-    left = text[:idx].strip()
-    right = text[idx:].strip()
-    return left, right
+    return text[:idx].strip(), text[idx:].strip()
 
 # -----------------------
-# The analyze function
+# Analyze Function
 # -----------------------
-def analyze(
-    albumin, creatinine, glucose, crp, mcv, rdw, alp,
-    wbc, lymph, age, gender, height, weight
-):
-    # Validate/constrain inputs
-    try:
-        age = int(age)
-    except Exception:
-        age = age
+def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
+            wbc, lymph, age, gender, height, weight):
+
+    # Calculate BMI
     try:
         height = float(height)
         weight = float(weight)
-        bmi = round(weight / ((height / 100) ** 2), 2) if height > 0 else "N/A"
+        bmi = round(weight / ((height / 100) ** 2), 2)
     except Exception:
         bmi = "N/A"
 
-    system_prompt = (
-    "You are 'Medical Insights AI', a trusted medical assistant.\n"
-    "You analyze patient demographics and biomarkers using two knowledge sources:\n"
-    "1. Internal medical expertise from pretraining.\n"
-    "2. The Knowledge Base of biomarker reference ranges provided below.\n\n"
-
-    "KNOWLEDGE BASE (STRICT PRIORITY):\n"
-    "- Albumin: 3.5 – 5.5 g/dL (Adults)\n"
-    "- Creatinine: 0.7 – 1.3 mg/dL (Adult Male), 0.6 – 1.1 mg/dL (Adult Female)\n"
-    "- Glucose (Fasting): 70 – 100 mg/dL (Adults)\n"
-    "- CRP: 0.3 – 10 mg/L (Adults)\n"
-    "- MCV: 80 – 100 fL (Adults)\n"
-    "- RDW: 11 – 15% (Adults)\n"
-    "- WBC: 4,000 – 11,000 /µL (Adults)\n"
-    "- Lymphocytes: 20 – 40% (Adults)\n"
-    "- ALP: 44 – 147 U/L (Adults)\n"
-
-    "RULES:\n"
-    "- For each biomarker, compare the patient’s value strictly against the ranges above.\n"
-    "- Always classify as Low / Normal / High.\n"
-    "- Provide 1–3 lines of interpretation (clinical meaning).\n"
-    "- Cite the source: Knowledge Base or Internal Knowledge.\n"
-    "- Do NOT invent biomarkers or repeat sections.\n"
-    "- Always integrate biomarker findings into system-level context.\n"
-    "- Provide actionable but non-prescriptive recommendations (nutrition, lifestyle, monitoring).\n\n"
-
-    "OUTPUT FORMAT (strict, mandatory):\n\n"
-
-    "*Executive Summary*\n"
-    "- Top Priority Issues: [List actual abnormalities, else 'None detected']\n"
-    "- Key Strengths: [Highlight normal/protective findings]\n\n"
-
-    "*Biomarker-by-Biomarker Analysis*\n"
-    "For each biomarker provided, follow this structure:\n"
-    "1. *Biomarker:* [Name]\n"
-    "   - *Patient Value:* [Value + Units]\n"
-    "   - *Reference Range:* [Range + Units + Population]\n"
-    "   - *Status:* [Low / Normal / High]\n"
-    "   - *Interpretation:* [Clear clinical meaning]\n"
-    "   - *Source:* [Knowledge Base / Internal Knowledge]\n\n"
-
-    "*System-Specific Analysis*\n"
-    "- Blood Health (MCV, RDW, WBC, Lymphocytes)\n"
-    "- Protein & Liver Health (Albumin, ALP)\n"
-    "- Kidney Health (Creatinine)\n"
-    "- Metabolic Health (Glucose, CRP)\n"
-    "- Anthropometrics (Age, Height, Weight, BMI)\n"
-    "- Other Systems: 'Not available from current biomarkers.'\n\n"
-
-    "*Personalized Action Plan*\n"
-    "- Medical: [Follow-up tests if abnormalities present]\n"
-    "- Nutrition: [Dietary guidance based on results, e.g., protein intake if albumin low, anti-inflammatory foods if CRP high]\n"
-    "- Lifestyle: [Exercise, hydration, stress, sleep tailored to BMI & biomarkers]\n"
-    "- Testing: [Mention ferritin, B12, folate, etc. ONLY if abnormalities suggest it]\n\n"
-
-    "*Interaction Alerts*\n"
-    "- Highlight important biomarker interactions (e.g., MCV+RDW for anemia, CRP+WBC for inflammation).\n\n"
-
-    "*Tabular Mapping*\n"
-    "- Markdown table with columns: | Biomarker |Reference range | Value | Status | AI-Inferred Insight |\n"
-    "- Include all available biomarkers in order.\n\n"
-
-    "*Enhanced AI Insights & Longitudinal Risk*\n"
-  
-
-    "STYLE:\n"
-    "- Professional, concise, medically accurate.\n"
-    "- Never output free-floating paragraphs, always follow structured format.\n"
-)
-
-
-
-
-
-
+    # Classify biomarkers via lookup
+    statuses = {
+        "Albumin": classify(albumin, "Albumin", gender),
+        "Creatinine": classify(creatinine, "Creatinine", gender),
+        "Glucose": classify(glucose, "Glucose", gender),
+        "CRP": classify(crp, "CRP", gender),
+        "MCV": classify(mcv, "MCV", gender),
+        "RDW": classify(rdw, "RDW", gender),
+        "WBC": classify(wbc * 1000, "WBC", gender),  # convert K/uL to /µL
+        "Lymphocytes": classify(lymph, "Lymphocytes", gender),
+        "ALP": classify(alp, "ALP", gender),
+    }
+
+    # Structured context for LLM
     patient_input = (
         f"Patient Profile:\n"
         f"- Age: {age}\n"
@@ -164,60 +105,46 @@ def analyze(
         f"- Height: {height} cm\n"
         f"- Weight: {weight} kg\n"
         f"- BMI: {bmi}\n\n"
-        "Biomarker Results:\n"
-        f"- Albumin: {albumin} g/dL\n"
-        f"- Creatinine: {creatinine} mg/dL\n"
-        f"- Glucose: {glucose} mg/dL\n"
-        f"- CRP: {crp} mg/L\n"
-        f"- MCV: {mcv} fL\n"
-        f"- RDW: {rdw} %\n"
-        f"- ALP: {alp} U/L\n"
-        f"- WBC: {wbc} K/uL\n"
-        f"- Lymphocytes: {lymph} %\n"
+        "Biomarker Results (Pre-classified):\n"
+    )
+    for biomarker, value in {
+        "Albumin": albumin,
+        "Creatinine": creatinine,
+        "Glucose": glucose,
+        "CRP": crp,
+        "MCV": mcv,
+        "RDW": rdw,
+        "ALP": alp,
+        "WBC": wbc,
+        "Lymphocytes": lymph,
+    }.items():
+        unit = REFERENCE_RANGES["Creatinine_Male"]["unit"] if biomarker == "Creatinine" else REFERENCE_RANGES[biomarker]["unit"]
+        patient_input += f"- {biomarker}: {value} {unit} → {statuses[biomarker]}\n"
+
+    system_prompt = (
+        "You are 'Medical Insights AI'.\n"
+        "The biomarker statuses (Low/Normal/High) are pre-computed using strict ranges.\n"
+        "Your role: generate detailed structured medical analysis, system-specific context, and recommendations.\n"
+        "Do NOT recalculate statuses — use the provided ones.\n"
+        "Follow the Executive Summary format strictly.\n"
     )
 
     prompt = system_prompt + "\n" + patient_input
 
-    # Generate
-    # Keep generation parameters conservative for Spaces
     gen = pipe(
         prompt,
-        max_new_tokens=2500,       # enough for executive summary
-        do_sample=False,          # greedy decoding (no randomness)
-        temperature=0.0,          # fully deterministic
-        top_p=1.0,                # consider full probability mass
-        repetition_penalty=1.0,   # neutral repetition handling
-        early_stopping=True,      # stop cleanly once done
-        return_full_text=False    # only assistant output
+        max_new_tokens=2000,
+        do_sample=False,
+        temperature=0.0,
+        return_full_text=False,
     )
 
-
-    # Extract generated text
-    generated = gen[0].get("generated_text") or gen[0].get("text") or str(gen[0])
-    generated = generated.strip()
-
-    # Clean: some models repeat prompt — attempt to strip prompt if present
-    # Remove leading prompt echo if it appears
-    if patient_input.strip() in generated:
-        generated = generated.split(patient_input.strip())[-1].strip()
-    # Also remove repeated instructions
-    if system_prompt.strip() in generated:
-        generated = generated.split(system_prompt.strip())[-1].strip()
-
-    # Split into left/right panels
-    left_md, right_md = split_report(generated)
-
-    # If the model output is empty or too short, return a helpful fallback
-    if len(left_md) < 50 and len(right_md) < 50:
-        fallback = (
-            "⚠️ The model returned an unexpectedly short response. Try re-running the report.\n\n"
-            "**Patient Profile:**\n" + patient_input
-        )
-        return fallback, ""
-    return left_md, right_md
+    generated = gen[0]["generated_text"].strip()
+    left, right = split_report(generated)
+    return left, right
 
 # -----------------------
-# Build Gradio app
+# Gradio UI (unchanged)
 # -----------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🏥 AI Medical Biomarker Dashboard")
@@ -261,7 +188,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         outputs=[left_output, right_output]
     )
 
-    
-# Launch (HF Spaces expects this pattern)
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+```