Update app.py

app.py

@@ -3,47 +3,50 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
 import torch
-import re
 
-
-MODEL_ID = "Muhammadidrees/my-biomed"
+# -----------------------
+# Recommended Biomedical Model (Swap here if needed)
+# -----------------------
+# Options:
+# - "stanford-crfm/BioMedLM"   (stable, PubMed-trained)
+# - "BioMistral/BioMistral-7B" (newer, PubMed + PMC heavy)
+# - "epfl-llm/ClinicalCamel"   (clinical reporting style)
+MODEL_ID = "Muhammadidrees/my-medgamma"
 
 # -----------------------
-# Load tokenizer + model safely (GPU or CPU)
+
+# Load tokenizer + model safely
 # -----------------------
 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,
+        device_map="auto",                  # auto GPU/CPU placement
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        low_cpu_mem_usage=True
+    )
+except Exception as e:
+    print(f"⚠️ GPU load failed, using CPU. Error: {e}")
+    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)
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=0 if torch.cuda.is_available() else -1
+)
 
 # -----------------------
-# Helper: robust section splitter
+# Helper: split report into panels
 # -----------------------
-def split_report(text):
-    """
-    Split model output into left (sections 1-4) and right (sections 5-6).
-    Accepts various markers for robustness.
-    """
-    # Normalize whitespace
+def split_report(text: str):
     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", "📊 Tabular", "## 5"]
     idx = None
     for m in markers:
         pos = text.find(m)
@@ -51,28 +54,16 @@ def split_report(text):
             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
-
+    return text[:idx].strip(), text[idx:].strip()
 
 # -----------------------
-# The analyze function
+# Main analysis 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:
@@ -84,206 +75,81 @@ def analyze(
     except Exception:
         bmi = "N/A"
 
-    # Reference Ranges (used for comparison)
-    reference_ranges = {
-        "Albumin": (3.5, 5.0),
-        "Creatinine": (0.6, 1.2),
-        "Glucose": (70, 99),
-        "CRP": (0, 3),
-        "MCV": (80, 100),
-        "RDW": (11.5, 14.5),
-        "ALP": (44, 147),
-        "WBC": (4.0, 11.0),
-        "Lymphocytes": (20, 40),
-    }
-
-    # Function to classify biomarker status
-    def classify_status(value, biomarker):
-        low, high = reference_ranges[biomarker]
-        if value < low:
-            return "Low"
-        elif value > high:
-            return "High"
-        else:
-            return "Normal"
-
-    # Biomarker status calculations
-    biomarkers = {
-        "Albumin": classify_status(float(albumin), "Albumin"),
-        "Creatinine": classify_status(float(creatinine), "Creatinine"),
-        "Glucose": classify_status(float(glucose), "Glucose"),
-        "CRP": classify_status(float(crp), "CRP"),
-        "MCV": classify_status(float(mcv), "MCV"),
-        "RDW": classify_status(float(rdw), "RDW"),
-        "ALP": classify_status(float(alp), "ALP"),
-        "WBC": classify_status(float(wbc), "WBC"),
-        "Lymphocytes": classify_status(float(lymph), "Lymphocytes"),
-    }
-
-    # Prepare the system prompt with biomarkers and patient info
+    # -----------------------
+    # Strict System Prompt
+    # -----------------------
     system_prompt = (
-    "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"
-        
-    "Reference Ranges (always use these as ground truth):\n"
-    "- Albumin: 3.5 – 5.0 g/dL\n"
-    "- Creatinine: 0.6 – 1.2 mg/dL\n"
-    "- Glucose (fasting): 70 – 99 mg/dL\n"
-    "- CRP: < 3 mg/L\n"
-    "- MCV: 80 – 100 fL\n"
-    "- RDW: 11.5 – 14.5 %\n"
-    "- ALP: 44 – 147 U/L\n"
-    "- WBC: 4.0 – 11.0 K/uL\n"
-    "- Lymphocytes: 20 – 40 %\n\n"
-    
-    "Check input of the user according to these ranges and provide analysis in a structured format.\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"
-    "- Status (Low/Normal/High) must be strictly defined 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"
-    "- 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 is 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"
-    "   - 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"
-    
-    "   - Each status column must be calculated by keeping the reference range in memory.\n"
-    "   - The table must contain the following structure:\n"
-    "     | Biomarker | Value |   Reference Range |Status (Low/Normal/High)| AI-Inferred Insight |\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"
+        "You are a professional AI Medical Assistant.\n"
+        "You must ONLY analyze: 9 Levine biomarkers + Age + Height + Weight.\n"
+        "Forbidden: Any extra labs (cholesterol, vitamin D, ferritin, ALT, AST, urine, hormones, genetics).\n"
+        "If information is not derivable, state clearly: 'Not available from current biomarkers.'\n\n"
+        "Biomarkers allowed:\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"
+        "Output format:\n"
+        "1. Executive Summary\n"
+        "2. System-Specific Analysis\n"
+        "3. Personalized Action Plan\n"
+        "4. Interaction Alerts\n"
+        "5. Tabular Mapping (Markdown table with Biomarker | Value | Range | Status | AI-Inferred Insight)\n"
+        "6. Enhanced AI Insights & Longitudinal Risk\n\n"
+        "Style: Professional, concise, structured, client-friendly. "
+        "No hallucinations. No extra biomarkers. No absolute longevity claims.\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"
+        f"Patient Profile:\n- Age: {age}\n- Gender: {gender}\n"
+        f"- Height: {height} cm\n- Weight: {weight} kg\n- BMI: {bmi}\n\n"
+        "Lab Values:\n"
+        f"- Albumin: {albumin}\n- Creatinine: {creatinine}\n"
+        f"- Glucose: {glucose}\n- CRP: {crp}\n"
+        f"- MCV: {mcv}\n- RDW: {rdw}\n"
+        f"- ALP: {alp}\n- WBC: {wbc}\n- Lymphocytes: {lymph}\n"
     )
 
-    # Prepare the final prompt
     prompt = system_prompt + "\n" + patient_input
 
-    # Generate the analysis
-    gen = pipe(prompt,
-               max_new_tokens=2500,
-               do_sample=True,
-               temperature=0.3,
-               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()
+    gen = pipe(
+        prompt,
+        max_new_tokens=1500,
+        do_sample=True,
+        temperature=0.25,  # lower temp = more factual
+        top_p=0.9,
+        repetition_penalty=1.05,
+        return_full_text=False
+    )
 
-    # 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()
+    generated = gen[0].get("generated_text", "").strip()
+    if not generated:
+        return "⚠️ No valid response. Please try again.", ""
 
-    # 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
+# 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):
-            gr.Markdown("### 👤 Demographics")
+        with gr.Column():
+            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)
 
-            gr.Markdown("### 🩸 Blood Panel")
+            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")
+        with gr.Column():
+            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)
@@ -293,12 +159,12 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             analyze_btn = gr.Button("🔬 Generate Report", variant="primary")
 
     with gr.Row():
-        with gr.Column(scale=1):
+        with gr.Column():
             gr.Markdown("### 📝 Summary & Action Plan")
-            left_output = gr.Markdown(value="Press *Generate Report* to create the analysis.")
-        with gr.Column(scale=1):
+            left_output = gr.Markdown()
+        with gr.Column():
             gr.Markdown("### 📊 Tabular & AI Insights")
-            right_output = gr.Markdown(value="Tabular mapping and enhanced insights will appear here.")
+            right_output = gr.Markdown()
 
     analyze_btn.click(
         fn=analyze,
@@ -306,7 +172,5 @@ 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)))