app.py

# app.py
import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import os
import torch
import re

MODEL_ID = "Muhammadidrees/BioMedLM-7B"

# -----------------------
# 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:
    # 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

# -----------------------
# The 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
    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_prompt = (
    "You are a professional AI Medical Assistant.\n"
    "You will analyze patient demographics and the Levine biomarker panel using strict reference ranges.\n\n"

    "REFERENCE RANGES (FOR INTERNAL USE — DO NOT REPEAT IN OUTPUT):\n"
    "- Albumin: 3.5 – 5.5 g/dL\n"
    "- Creatinine: 0.6 – 1.3 mg/dL\n"
    "- Glucose (fasting): 70 – 99 mg/dL\n"
    "- CRP: < 3.0 mg/L (low risk, <1 optimal)\n"
    "- MCV: 80 – 100 fL\n"
    "- RDW: 11 – 15 %\n"
    "- ALP: 20 – 120 U/L\n"
    "- WBC: 4 – 11 K/uL\n"
    "- Lymphocytes: 20 – 45 %\n\n"

    "RULES:\n"
    "- Classify each biomarker strictly as Low / Normal / High using the ranges above.\n"
    "- Do NOT output or repeat the ranges themselves.\n"
    "- If all biomarkers are normal: explicitly state 'No abnormalities detected.'\n"
    "- If a section has no relevant data: state 'Not available from current biomarkers.'\n"
    "- Only mention follow-up tests if abnormalities are detected.\n"
    "- Keep all analysis range-based, avoid speculation or invented risks.\n\n"

    "OUTPUT FORMAT:\n"
    "1. Executive Summary\n"
    "   - Top Priority Issues (or 'None detected')\n"
    "   - Key Strengths\n\n"
    "2. 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"
    "3. Personalized Action Plan\n"
    "   - Medical\n"
    "   - Nutrition\n"
    "   - Lifestyle\n"
    "   - Testing (only if issues detected)\n\n"
    "4. Interaction Alerts\n"
    "   - Highlight only valid cross-biomarker relationships.\n\n"
    "5. Tabular Mapping\n"
    "   - Markdown table with 4 columns: | Biomarker | Value | Status | AI-Inferred Insight |\n"
    "   - Include all 9 biomarkers, in this order: Albumin, Creatinine, Glucose, CRP, MCV, RDW, ALP, WBC, Lymphocytes.\n\n"
    "6. Enhanced AI Insights & Longitudinal Risk\n"
    "   - Only if supported by patterns, else state 'Not available from current biomarkers.'\n\n"

    "STYLE:\n"
    "- Professional, concise, medically accurate.\n"
    "- No hallucinations. No extra biomarkers. No repetition of reference ranges.\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"
        "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"
    )

    prompt = system_prompt + "\n" + patient_input

    # 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])
    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):
            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")
            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)
            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.")

    analyze_btn.click(
        fn=analyze,
        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp, wbc, lymph, age, gender, height, weight],
        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)))