Update app.py

app.py

@@ -1,115 +1,256 @@
+# app.py
 import gradio as gr
-import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import os
+import torch
+import re
 
-MODEL_ID = "Muhammadidrees/bioLLM"
+MODEL_ID = "Muhammadidrees/BioLLM"
 
-#import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+# -----------------------
+# Load tokenizer + model safely (GPU or CPU)
+# -----------------------
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-# Load your model from Hugging Face Hub
-# MODEL_ID = "Muhammadidrees/MedicalInsights"
+# 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)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
+# 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
 
-# Function to build structured input and query the LLM
+    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
 ):
-    # Calculate BMI (hidden from user, only passed to LLM)
+    # Validate/constrain inputs
     try:
-        height_m = height / 100  # cm → m
-        bmi = round(weight / (height_m ** 2), 2)
+        age = int(age)
+    except Exception:
+        age = age
+    try:
+        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-style instruction
     system_prompt = (
-        "You are an advanced AI medical assistant. "
-        "Analyze the patient’s biomarkers and demographics. "
-        "Provide a structured assessment including: "
-         """Executive Summary,System-Specific Analysis,Personalized Action Plan,Interaction Alerts,Longevity Metrics,Enhanced AI Insights & Longitudinal Risk
-       
-"""    )
-
-    # Construct patient profile input
-    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
-    - C-Reactive Protein: {crp} mg/L
-    - Mean Cell Volume: {mcv} fL
-    - Red Cell Distribution Width: {rdw} %
-    - Alkaline Phosphatase: {alp} U/L
-    - White Blood Cell Count: {wbc} K/uL
-    - Lymphocyte Percentage: {lymph} %
-    """
+    "You are a professional AI Medical Assistant.\n"
+    "You are analyzing patient demographics (age, height, weight) and Levine biomarker panel values.\n\n"
 
-    prompt = system_prompt + "\n" + patient_input
+    "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"
+    "- 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)\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"
 
-    # Call LLM (exclude echoed input)
-    result = pipe(
-        prompt,
-        max_new_tokens=2000,
-        do_sample=True,
-        temperature=0.6,
-        return_full_text=False
+    "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 four columns:\n"
+    "     | Biomarker | Value | Status (Low/Normal/High) | AI-Inferred Insight |\n"
+    "   - Include ALL 9 Levine biomarkers (Albumin, Creatinine, Glucose, CRP, MCV, RDW, ALP, WBC, Lymphocytes).\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:\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 result[0]["generated_text"].strip()
 
+    prompt = system_prompt + "\n" + patient_input
 
-# Build Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧪 Medical Insights AI — Enter Patient Data")
+    # 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)
 
-    with gr.Row():
-        albumin = gr.Number(label="Albumin (g/dL)")
-        wbc = gr.Number(label="White Blood Cell Count (K/uL)")
+    # Extract generated text
+    generated = gen[0].get("generated_text") or gen[0].get("text") or str(gen[0])
+    generated = generated.strip()
 
-    with gr.Row():
-        creatinine = gr.Number(label="Creatinine (mg/dL)")
-        lymph = gr.Number(label="Lymphocyte Percentage (%)")
+    # 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()
 
-    with gr.Row():
-        glucose = gr.Number(label="Glucose (mg/dL)")
-        age = gr.Number(label="Age (years)")
+    # Split into left/right panels
+    left_md, right_md = split_report(generated)
 
-    with gr.Row():
-        crp = gr.Number(label="C-Reactive Protein (mg/L)")
-        gender = gr.Dropdown(choices=["Male", "Female"], label="Gender")
+    # 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
 
-    with gr.Row():
-        mcv = gr.Number(label="Mean Cell Volume (fL)")
-        height = gr.Number(label="Height (cm)")
+# -----------------------
+# 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():
-        rdw = gr.Number(label="Red Cell Distribution Width (%)")
-        weight = gr.Number(label="Weight (kg)")
+        with gr.Column(scale=1):
+            gr.Markdown("### 👤 Demographics")
+            age = gr.Number(label="Age", value=45)
+            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)
 
-    with gr.Row():
-        alp = gr.Number(label="Alkaline Phosphatase (U/L)")
+            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)
 
-    analyze_btn = gr.Button("🔎 Analyze")
-    output = gr.Textbox(label="AI Medical Assessment", lines=12)
+        with gr.Column(scale=1):
+            gr.Markdown("### 🧬 Chemistry Panel")
+            albumin = gr.Number(label="Albumin (g/dL)", value=4.2)
+            creatinine = gr.Number(label="Creatinine (mg/dL)", value=0.9)
+            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(value="Press *Generate Report* to create the analysis.")
+        with gr.Column(scale=1):
+            gr.Markdown("### 📊 Tabular & AI Insights")
+            right_output = gr.Markdown(value="Tabular mapping and enhanced insights will appear here.")
 
-    # Run analysis
     analyze_btn.click(
         fn=analyze,
-        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
-                wbc, lymph, age, gender, height, weight],
-        outputs=output
+        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc, lymph, age, gender, height, weight],
+        outputs=[left_output, right_output]
     )
 
-demo.launch()
+    
+# 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)))