Update app.py

app.py

@@ -1,120 +1,178 @@
+# app.py
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import os
+import torch
+import re
 
-# -----------------------
-# Load Hugging Face Model
-# -----------------------
 MODEL_ID = "Muhammadidrees/my-gpt-oss"
 
+# -----------------------
+# Load tokenizer + model safely (GPU or CPU)
+# -----------------------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID)  # works on CPU & GPU
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
+# 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)
+
+# Create pipeline
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
+
+# -----------------------
+# Helper: robust section splitter
+# -----------------------
+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
+    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
 
 # -----------------------
-# Analysis Function
+# The analyze function
 # -----------------------
 def analyze(
     albumin, creatinine, glucose, crp, mcv, rdw, alp,
     wbc, lymph, age, gender, height, weight
 ):
-    # Calculate BMI
+    # Validate/constrain inputs
+    try:
+        age = int(age)
+    except Exception:
+        age = age
     try:
-        height_m = height / 100
-        bmi = round(weight / (height_m ** 2), 2)
+        height = float(height)
+        weight = float(weight)
+        bmi = round(weight / ((height / 100) ** 2), 2) if height > 0 else "N/A"
     except Exception:
         bmi = "N/A"
 
-    system_prompt = """
-You are a professional AI Medical Assistant.
-You are analyzing patient demographics and Levine biomarker panel values.
-Output MUST strictly follow this structured format:
-
-1. Executive Summary
-   - Top Priority Issues
-   - Key Strengths
-
-2. System-Specific Analysis
-   - Blood Health (MCV, RDW, Lymphocytes, WBC)
-   - Protein & Liver Health (Albumin, ALP)
-   - Kidney Health (Creatinine µmol/L)
-   - Metabolic Health (Glucose mmol/L, CRP)
-   - Other relevant systems
-
-3. Personalized Action Plan
-   - Medical (tests/consults)
-   - Nutrition (diet & supplements)
-   - Lifestyle (hydration, exercise, sleep)
-   - Testing (follow-up labs: ferritin, Vitamin D, GGT)
-
-4. Interaction Alerts
-   - How biomarkers interact (e.g., anemia ↔ infection cycle, ALP with bone/liver origin)
-
-5. Tabular Mapping (Biomarker → Value → Status → AI-Inferred Insight → Client-Friendly Message)
-
-6. Enhanced AI Insights & Longitudinal Risk
-   - Subclinical nutrient predictions (Iron, B12, Folate, Copper)
-   - Elevated ALP interpretation (bone vs liver origin)
-   - WBC & lymphocyte trends for immunity
-   - Predictive longevity risk profile
-"""
-
-    patient_input = f"""
-Patient Profile:
-- Age: {age}
-- Gender: {gender}
-- Height: {height} cm
-- Weight: {weight} kg
-- BMI: {bmi}
-
-Lab Values:
-- Albumin: {albumin} g/dL
-- Creatinine: {creatinine} mg/dL
-- Glucose: {glucose} mg/dL
-- CRP: {crp} mg/L
-- MCV: {mcv} fL
-- RDW: {rdw} %
-- ALP: {alp} U/L
-- WBC: {wbc} K/uL
-- Lymphocytes: {lymph} %
-"""
-
-    prompt = system_prompt + "\n" + patient_input
-
-    result = pipe(
-        prompt,
-        max_new_tokens=1000,
-        do_sample=True,
-        temperature=0.3,
-        top_p=0.9
+    # Structured system prompt (no estimates; follow sections exactly)
+    system_prompt = (
+        "You are a professional AI Medical Assistant.\n"
+        "You are analyzing patient demographics and Levine biomarker panel values.\n"
+        "Output MUST strictly follow this structured format (no extra commentary):\n\n"
+        "1. Executive Summary\n"
+        "   - Top Priority Issues\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 µmol/L)\n"
+        "   - Metabolic Health (Glucose mmol/L, CRP)\n"
+        "   - Other relevant systems\n\n"
+        "3. Personalized Action Plan\n"
+        "   - Medical (tests/consults)\n"
+        "   - Nutrition (diet & supplements)\n"
+        "   - Lifestyle (hydration, exercise, sleep)\n"
+        "   - Testing (follow-up labs: ferritin, Vitamin D, GGT)\n\n"
+        "4. Interaction Alerts\n"
+        "   - How biomarkers interact (e.g., anemia ↔ infection cycle, ALP with bone/liver origin)\n\n"
+        "5. Tabular Mapping (Biomarker → Value → Status → AI-Inferred Insight → Client-Friendly Message)\n\n"
+        "6. Enhanced AI Insights & Longitudinal Risk\n"
+        "   - Subclinical nutrient predictions (Iron, B12, Folate, Copper)\n"
+        "   - Elevated ALP interpretation (bone vs liver origin)\n"
+        "   - WBC & lymphocyte trends for immunity\n"
+        "   - Predictive longevity risk profile\n\n"
+        "Now analyze the patient below and produce the report strictly in the format above using bullet points, headings and a Markdown table for section 5.\n"
     )
 
-    text = result[0]["generated_text"].strip()
-
-    # Robust split into left (1–4) and right (5–6)
-    left_sections, right_sections = [], []
-    capture_left = True
-    for line in text.splitlines():
-        if "Tabular Mapping" in line:
-            capture_left = False
-        if capture_left:
-            left_sections.append(line)
-        else:
-            right_sections.append(line)
+    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:\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"
+    )
 
-    return "\n".join(left_sections), "\n".join(right_sections)
+    prompt = system_prompt + "\n" + patient_input
 
+    # Generate
+    # Keep generation parameters conservative for Spaces
+    gen = pipe(prompt,
+               max_new_tokens=900,
+               do_sample=True,
+               temperature=0.25,
+               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])
+    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
 
 # -----------------------
-# Gradio Interface
+# Build Gradio app
 # -----------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    
-    gr.Markdown("""
-    # 🏥 AI Medical Biomarker Analysis  
-    ### Comprehensive wellness insights from lab values
-    """)
-    
+    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):
             gr.Markdown("### 👤 Demographics")
@@ -122,13 +180,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             gender = gr.Dropdown(["Male", "Female"], label="Gender", value="Male")
             height = gr.Number(label="Height (cm)", value=174)
             weight = gr.Number(label="Weight (kg)", value=75)
-            
+
             gr.Markdown("### 🩸 Blood Panel")
             wbc = gr.Number(label="WBC (K/uL)", value=6.5)
             lymph = gr.Number(label="Lymphocytes (%)", value=30)
             mcv = gr.Number(label="MCV (fL)", value=88)
             rdw = gr.Number(label="RDW (%)", value=13)
-        
+
         with gr.Column(scale=1):
             gr.Markdown("### 🧬 Chemistry Panel")
             albumin = gr.Number(label="Albumin (g/dL)", value=4.2)
@@ -136,28 +194,25 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             glucose = gr.Number(label="Glucose (mg/dL)", value=92)
             crp = gr.Number(label="CRP (mg/L)", value=1.0)
             alp = gr.Number(label="ALP (U/L)", value=70)
-            
+
             analyze_btn = gr.Button("🔬 Generate Report", variant="primary")
-    
+
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 📝 Summary & Action Plan")
-            left_output = gr.Markdown()
+            left_output = gr.Markdown(value="Press *Generate Report* to create the analysis.")
         with gr.Column(scale=1):
             gr.Markdown("### 📊 Tabular & AI Insights")
-            right_output = gr.Markdown()
-    
+            right_output = gr.Markdown(value="Tabular mapping and enhanced insights will appear here.")
+
     analyze_btn.click(
         fn=analyze,
-        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
-                wbc, lymph, age, gender, height, weight],
+        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc, lymph, age, gender, height, weight],
         outputs=[left_output, right_output]
     )
-    
-    gr.Markdown(
-        "*⚠️ Disclaimer: This AI output is for educational purposes only and not a substitute for professional medical advice.*"
-    )
 
+    gr.Markdown("*⚠️ Disclaimer: This AI output is for educational purposes only and not a substitute for professional medical advice.*")
 
+# Launch (HF Spaces expects this pattern)
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))