Upload app.py

app.py

@@ -3,80 +3,78 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
 import torch
+import re
 
 MODEL_ID = "Muhammadidrees/MedicalInsights"
 
 # -----------------------
-# Load tokenizer + model
+# Load tokenizer + model safely (GPU or CPU)
 # -----------------------
 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:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True
-    )
+    # Fallback: force CPU (slower but safe)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True)
 
-device = 0 if torch.cuda.is_available() else -1
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
+# Create pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
 # -----------------------
-# Reference ranges lookup
+# Helper: robust section splitter
 # -----------------------
-REFERENCE_RANGES = {
-    "Albumin": (3.5, 5.5),
-    "Creatinine": {"Male": (0.7, 1.3), "Female": (0.6, 1.1)},
-    "Glucose": (70, 100),
-    "CRP": (0.3, 10),
-    "MCV": (80, 100),
-    "RDW": (11, 15),
-    "WBC": (4, 11),  # K/uL
-    "Lymphocytes": (20, 40),
-    "ALP": (44, 147),
-}
-
-def classify_biomarker(name, value, gender=None):
+def split_report(text):
     """
-    Returns Low / Normal / High for a biomarker using REFERENCE_RANGES.
-    Safe: won't crash if value or range missing.
+    Split model output into left (sections 1-4) and right (sections 5-6).
+    Accepts various markers for robustness.
     """
-    try:
-        if value is None:
-            return "Unknown"
-        ref = REFERENCE_RANGES.get(name)
-        if ref is None:
-            return "Unknown"
-
-        if isinstance(ref, dict) and gender:
-            low, high = ref.get(gender, (None, None))
-        else:
-            low, high = ref
-
-        if low is None or high is None:
-            return "Unknown"
-
-        if value < low:
-            return "Low"
-        elif value > high:
-            return "High"
-        else:
-            return "Normal"
-    except Exception as e:
-        return f"Error: {str(e)}"
+    # 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
+    idx = None
+    for m in markers:
+        pos = text.find(m)
+        if pos != -1:
+            if 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
 
 # -----------------------
 # The analyze function
 # -----------------------
-def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
-            wbc, lymph, age, gender, height, weight):
-
-    # Safe conversions
+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 = "N/A"
-
+        age = age
     try:
         height = float(height)
         weight = float(weight)
@@ -84,86 +82,139 @@ def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
     except Exception:
         bmi = "N/A"
 
-    # Classify each biomarker
-    statuses = {
-        "Albumin": classify_biomarker("Albumin", albumin, gender),
-        "Creatinine": classify_biomarker("Creatinine", creatinine, gender),
-        "Glucose": classify_biomarker("Glucose", glucose),
-        "CRP": classify_biomarker("CRP", crp),
-        "MCV": classify_biomarker("MCV", mcv),
-        "RDW": classify_biomarker("RDW", rdw),
-        "WBC": classify_biomarker("WBC", wbc),
-        "Lymphocytes": classify_biomarker("Lymphocytes", lymph),
-        "ALP": classify_biomarker("ALP", alp),
-    }
-
-    # Build structured patient input for LLM
-    patient_input = (
-        f"Patient Profile:\n"
-        f"- Age: {age}\n"
-        f"- Gender: {gender}\n"
-        f"- Height: {height} cm\n"
-        f"- Weight: {weight} kg\n"
-        f"- BMI: {bmi}\n\n"
-        "Biomarker Results:\n"
-    )
-    for biomarker, value in {
-        "Albumin": albumin,
-        "Creatinine": creatinine,
-        "Glucose": glucose,
-        "CRP": crp,
-        "MCV": mcv,
-        "RDW": rdw,
-        "ALP": alp,
-        "WBC": wbc,
-        "Lymphocytes": lymph,
-    }.items():
-        patient_input += f"- {biomarker}: {value} ({statuses[biomarker]})\n"
-
-    # System prompt
     system_prompt = (
-        "You are 'Medical Insights AI', a trusted medical assistant.\n"
-        "The biomarker classifications (Low / Normal / High) are already provided.\n"
-        "Do not recompute them — just use them to generate:\n\n"
-        "*Executive Summary*\n"
-        "- Top Priority Issues\n"
-        "- Key Strengths\n\n"
-        "*System-Specific Analysis*\n"
-        "- Blood Health\n"
-        "- Protein & Liver Health\n"
-        "- Kidney Health\n"
-        "- Metabolic Health\n"
-        "- Anthropometrics\n"
-        "- Other Systems\n\n"
-        "*Personalized Action Plan*\n"
-        "- Medical, Nutrition, Lifestyle, Testing\n\n"
-        "*Interaction Alerts*\n\n"
-        "*Tabular Mapping*\n"
-        "- Table with Biomarker | Reference Range | Value | Status | Insight\n\n"
-        "*Enhanced AI Insights & Longitudinal Risk*\n"
-    )
+    "You are a professional AI Medical Assistant.\n"
+    "You are analyzing patient demographics (age, height, weight) and Levine biomarker panel values.\n\n"
+
+    "The Levine biomarker panel includes:\n"
+    "- Albumin\n"
+    "- Creatinine\n"
+    "- Glucose\n"
+    "- C-reactive protein (CRP)\n"
+    "- Mean Cell Volume (MCV)\n"
+    "- Red Cell Distribution Width (RDW)\n"
+    "- Alkaline Phosphatase (ALP)\n"
+    "- White Blood Cell count (WBC)\n"
+    "- Lymphocyte percentage\n\n"
+
+    "STRICT RULES:\n"
+    "- Use ONLY the 9 biomarkers above + age, height, weight.\n"
+    "- DO NOT use or invent other lab results (e.g., cholesterol, vitamin D, ferritin, ALT, AST, urine results).\n"
+    "- Use the reference ranges provided in the Patient Summary for all biomarker interpretations.\n"
+    "- Status (Low/Normal/High) must be defined strictly by whether the value falls below, within, or above the provided reference ranges.\n"
+    "- Analysis must always explain deviations with respect to these reference ranges.\n"
+    "- If a section cannot be addressed with available data, explicitly state: 'Not available from current biomarkers.'\n"
+    "- Do not give absolute longevity scores. Instead, summarize trends (e.g., 'No major abnormalities suggesting elevated short-term risk.').\n"
+    "- Nutrient status (Iron, B12, Folate) can only be suggested as possible IF supported by MCV + RDW patterns, but never stated as confirmed.\n"
+    "- Interpret ALP cautiously: mention bone vs liver as possible sources, but highlight that more tests would be required to confirm.\n"
+    "- Always highlight limitations where applicable.\n\n"
+
+    "OUTPUT FORMAT (strict, structured, and professional):\n\n"
+
+    "1. Executive Summary\n"
+    "   - Top Priority Issues (based only on provided biomarkers and their ranges)\n"
+    "   - Key Strengths\n\n"
+
+    "2. System-Specific Analysis\n"
+    "   - Blood Health (MCV, RDW, Lymphocytes, WBC)\n"
+    "   - Protein & Liver Health (Albumin, ALP)\n"
+    "   - Kidney Health (Creatinine)\n"
+    "   - Metabolic Health (Glucose, CRP)\n"
+    "   - Anthropometrics (Age, Height, Weight, BMI)\n"
+    "   - Other systems: Always state 'Not available from current biomarkers.' if data missing\n\n"
+
+    "3. Personalized Action Plan\n"
+    "   - Medical (tests/consults related only to biomarkers — e.g., repeat CBC, iron studies if anemia suspected)\n"
+    "   - Nutrition (diet & supplements grounded ONLY in biomarker findings — e.g., protein intake if albumin low, anti-inflammatory foods if CRP elevated)\n"
+    "   - Lifestyle (hydration, exercise, sleep — general guidance contextualized by BMI and biomarkers)\n"
+    "   - Testing (only mention ferritin, B12, folate, GGT, etc. as follow-up — but clarify these are NOT part of current data)\n\n"
+
+    "4. Interaction Alerts\n"
+    "   - Describe ONLY interactions among provided biomarkers (e.g., RDW with MCV for anemia trends, ALP bone/liver origin, WBC with CRP for infection/inflammation)\n\n"
+
+    "5. Tabular Mapping\n"
+    "   - This section must always include a Markdown table.\n"
+    "   - The table must contain exactly five columns:\n"
+    "     | Biomarker | Value | Status (Low/Normal/High) | Reference Range | AI-Inferred Insight |\n"
+    "   - The Reference Range column must be populated with the ranges given in the Patient Summary.\n"
+    "   - The first row after the header must begin directly with 'Albumin'.\n"
+    "   - Do NOT add any index numbers (0,1,2...) or empty rows.\n"
+    "   - Each biomarker must appear exactly once as a separate row.\n\n"
+
+    "6. Enhanced AI Insights & Longitudinal Risk\n"
+    "   - Subclinical nutrient predictions ONLY if patterns (MCV + RDW) suggest it — state as possible, not confirmed.\n"
+    "   - ALP interpretation limited to bone vs liver origin (uncertain without further tests).\n"
+    "   - WBC & lymphocyte balance for immunity.\n"
+    "   - Risk framing: Highlight if biomarkers suggest resilience or potential stress, but avoid absolute longevity claims.\n\n"
+
+    "STYLE REQUIREMENTS:\n"
+    "- Use clear section headings and bullet points.\n"
+    "- Keep language professional, concise, and client-friendly.\n"
+    "- Format tables cleanly in Markdown.\n"
+    "- Present output beautifully, like a polished medical summary.\n"
+)
+
+
+    patient_input =(
+    f"Patient Profile:\n"
+    f"- Age: {age}\n"
+    f"- Gender: {gender}\n"
+    f"- Height: {height} cm\n"
+    f"- Weight: {weight} kg\n"
+    f"- BMI: {bmi}\n\n"
+    "Lab Values (with Normal Ranges):\n"
+    f"- Albumin: {albumin} g/dL (Normal: 3.5 – 5.0 g/dL)\n"
+    f"- Creatinine: {creatinine} mg/dL (Normal: 0.6 – 1.2 mg/dL)\n"
+    f"- Glucose: {glucose} mg/dL (Normal: 70 – 99 mg/dL, fasting)\n"
+    f"- CRP: {crp} mg/L (Normal: < 3 mg/L)\n"
+    f"- MCV: {mcv} fL (Normal: 80 – 100 fL)\n"
+    f"- RDW: {rdw} % (Normal: 11.5 – 14.5 %)\n"
+    f"- ALP: {alp} U/L (Normal: 44 – 147 U/L)\n"
+    f"- WBC: {wbc} K/uL (Normal: 4.0 – 11.0 K/uL)\n"
+    f"- Lymphocytes: {lymph} % (Normal: 20 – 40 %)\n"
+)
 
     prompt = system_prompt + "\n" + patient_input
 
-    gen = pipe(
-        prompt,
-        max_new_tokens=1200,
-        do_sample=False,
-        temperature=0.0,
-        top_p=1.0,
-        repetition_penalty=1.0,
-        early_stopping=True,
-        return_full_text=False
-    )
+    # Generate
+    # Keep generation parameters conservative for Spaces
+    gen = pipe(prompt,
+               max_new_tokens=2500,
+               do_sample=True,
+               temperature=0.001,
+               top_p=0.9,
+               return_full_text=False)
 
+    # Extract generated text
     generated = gen[0].get("generated_text") or gen[0].get("text") or str(gen[0])
-    return generated, generated  # left and right panels (split optional)
+    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
 
 # -----------------------
 # Build Gradio app
 # -----------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🏥 AI Medical Biomarker Dashboard")
+    gr.Markdown("Enter lab values and demographics — Report is generated in two panels (Summary & Table/Insights).")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -203,10 +254,7 @@ 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)),
-        show_error=True,   # Show traceback on Hugging Face
-        debug=True
-    )
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))