Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import os
@@ -6,7 +7,7 @@ import torch
 MODEL_ID = "Muhammadidrees/MedicalInsights"
 
 # -----------------------
-# Load tokenizer + model safely (GPU or CPU)
+# Load tokenizer + model
 # -----------------------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
@@ -17,99 +18,86 @@ except Exception:
         MODEL_ID, torch_dtype=torch.float32, low_cpu_mem_usage=True
     )
 
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1,
-)
+device = 0 if torch.cuda.is_available() else -1
+pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
 # -----------------------
-# Lookup Table for Biomarkers
+# Reference ranges lookup
 # -----------------------
 REFERENCE_RANGES = {
-    "Albumin": {"low": 3.5, "high": 5.5, "unit": "g/dL"},
-    "Creatinine_Male": {"low": 0.7, "high": 1.3, "unit": "mg/dL"},
-    "Creatinine_Female": {"low": 0.6, "high": 1.1, "unit": "mg/dL"},
-    "Glucose": {"low": 70, "high": 100, "unit": "mg/dL"},
-    "CRP": {"low": 0.3, "high": 10, "unit": "mg/L"},
-    "MCV": {"low": 80, "high": 100, "unit": "fL"},
-    "RDW": {"low": 11, "high": 15, "unit": "%"},
-    "WBC": {"low": 4000, "high": 11000, "unit": "/µL"},
-    "Lymphocytes": {"low": 20, "high": 40, "unit": "%"},
-    "ALP": {"low": 44, "high": 147, "unit": "U/L"},
+    "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(value, biomarker, gender="Male"):
-    """Classify biomarker as Low, Normal, or High using lookup table."""
-    if biomarker == "Creatinine":
-        ref = REFERENCE_RANGES[f"Creatinine_{gender}"]
-    else:
-        ref = REFERENCE_RANGES[biomarker]
-
-    if value < ref["low"]:
-        return "Low"
-    elif value > ref["high"]:
-        return "High"
-    else:
-        return "Normal"
-
-# -----------------------
-# Splitter helper
-# -----------------------
-def split_report(text):
-    text = text.strip()
-    markers = ["5. Tabular Mapping", "Tabular Mapping", "📊 Tabular", "## 5"]
-    idx = None
-    for m in markers:
-        pos = text.find(m)
-        if pos != -1 and (idx is None or pos < idx):
-            idx = pos
-    if idx is None:
-        return text, ""
-    return text[:idx].strip(), text[idx:].strip()
-
-# -----------------------
-# Post-processing cleaner (remove repeated lines)
-# -----------------------
-def clean_output(text):
-    lines = text.split("\n")
-    cleaned = []
-    prev_line = None
-    for line in lines:
-        if line.strip() and line != prev_line:  # remove duplicates
-            cleaned.append(line)
-        prev_line = line
-    return "\n".join(cleaned)
+def classify_biomarker(name, value, gender=None):
+    """
+    Returns Low / Normal / High for a biomarker using REFERENCE_RANGES.
+    Safe: won't crash if value or range missing.
+    """
+    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)}"
 
 # -----------------------
-# Analyze Function
+# The analyze function
 # -----------------------
 def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
             wbc, lymph, age, gender, height, weight):
 
-    # Calculate BMI
+    # Safe conversions
+    try:
+        age = int(age)
+    except Exception:
+        age = "N/A"
+
     try:
         height = float(height)
         weight = float(weight)
-        bmi = round(weight / ((height / 100) ** 2), 2)
+        bmi = round(weight / ((height / 100) ** 2), 2) if height > 0 else "N/A"
     except Exception:
         bmi = "N/A"
 
-    # Classify biomarkers via lookup
+    # Classify each biomarker
     statuses = {
-        "Albumin": classify(albumin, "Albumin", gender),
-        "Creatinine": classify(creatinine, "Creatinine", gender),
-        "Glucose": classify(glucose, "Glucose", gender),
-        "CRP": classify(crp, "CRP", gender),
-        "MCV": classify(mcv, "MCV", gender),
-        "RDW": classify(rdw, "RDW", gender),
-        "WBC": classify(wbc * 1000, "WBC", gender),  # convert K/uL to /µL
-        "Lymphocytes": classify(lymph, "Lymphocytes", gender),
-        "ALP": classify(alp, "ALP", gender),
+        "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),
     }
 
-    # Structured context for LLM
+    # Build structured patient input for LLM
     patient_input = (
         f"Patient Profile:\n"
         f"- Age: {age}\n"
@@ -117,7 +105,7 @@ def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
         f"- Height: {height} cm\n"
         f"- Weight: {weight} kg\n"
         f"- BMI: {bmi}\n\n"
-        "Biomarker Results (Pre-classified):\n"
+        "Biomarker Results:\n"
     )
     for biomarker, value in {
         "Albumin": albumin,
@@ -130,49 +118,52 @@ def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
         "WBC": wbc,
         "Lymphocytes": lymph,
     }.items():
-        unit = REFERENCE_RANGES["Creatinine_Male"]["unit"] if biomarker == "Creatinine" else REFERENCE_RANGES[biomarker]["unit"]
-        patient_input += f"- {biomarker}: {value} {unit} → {statuses[biomarker]}\n"
+        patient_input += f"- {biomarker}: {value} ({statuses[biomarker]})\n"
 
+    # System prompt
     system_prompt = (
-        "You are a professional AI Medical Assistant.\n"
-        "Generate a structured medical report ONLY using the 9 Levine biomarkers, age, height, and weight.\n"
-        "Follow this structure:\n\n"
-        "1. Executive Summary\n"
-        "2. System-Specific Analysis\n"
-        "3. Personalized Action Plan\n"
-        "4. Interaction Alerts\n"
-        "5. Tabular Mapping (with 4 columns: Biomarker | Value | Status | AI-Inferred Insight)\n"
-        "6. Enhanced AI Insights & Longitudinal Risk\n\n"
-        "Rules:\n"
-        "- Always include all 9 biomarkers in the table.\n"
-        "- Do not repeat lines or add unrelated labs.\n"
-        "- Keep responses concise, professional, and readable.\n"
+        "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"
     )
 
     prompt = system_prompt + "\n" + patient_input
 
-    # ---- FIXED PIPE CALL ----
     gen = pipe(
         prompt,
         max_new_tokens=1200,
-        do_sample=True,          # controlled sampling
-        temperature=0.3,
-        top_p=0.9,
-        repetition_penalty=1.2,  # prevent loops
-        eos_token_id=tokenizer.eos_token_id,
-        return_full_text=False,
+        do_sample=False,
+        temperature=0.0,
+        top_p=1.0,
+        repetition_penalty=1.0,
+        early_stopping=True,
+        return_full_text=False
     )
 
-    generated = clean_output(gen[0]["generated_text"].strip())
-    left, right = split_report(generated)
-    return left, right
+    generated = gen[0].get("generated_text") or gen[0].get("text") or str(gen[0])
+    return generated, generated  # left and right panels (split optional)
 
 # -----------------------
-# Gradio UI
+# 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):
@@ -213,4 +204,9 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.environ.get("PORT", 7860)))
+    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
+    )