Update app.py

app.py

@@ -1,100 +1,76 @@
-# app.py
-import gradio as gr
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import gradio as gr
 import os
-import torch
 
-# -----------------------
-# 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"
+# ----------------------------
+# Model Config
+# ----------------------------
+MODEL_ID = "Muhammadidrees/my-biomed"   # or "BioMistral/BioMistral-7B"
 
-# -----------------------
-
-# Load tokenizer + model safely
-# -----------------------
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
-try:
-    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
-    )
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",           # Ensures GPU usage on HF (L4)
+    torch_dtype=torch.float16    # FP16 for speed + memory efficiency
+)
 
 pipe = pipeline(
     "text-generation",
     model=model,
-    tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1
+    tokenizer=tokenizer
 )
 
-# -----------------------
+# ----------------------------
 # Helper: split report into panels
-# -----------------------
+# ----------------------------
 def split_report(text: str):
-    text = text.strip()
     markers = ["5. Tabular", "📊 Tabular", "## 5"]
     idx = None
     for m in markers:
         pos = text.find(m)
         if pos != -1:
-            if idx is None or pos < idx:
-                idx = pos
+            idx = pos if idx is None or pos < idx else idx
     if idx is None:
-        return text, ""
+        return text.strip(), ""
     return text[:idx].strip(), text[idx:].strip()
 
-# -----------------------
-# Main analysis function
-# -----------------------
-def analyze(
-    albumin, creatinine, glucose, crp, mcv, rdw, alp,
-    wbc, lymph, age, gender, height, weight
-):
-    try:
-        age = int(age)
-    except Exception:
-        age = age
+# ----------------------------
+# Main Analysis Function
+# ----------------------------
+def analyze(albumin, creatinine, glucose, crp, mcv, rdw, alp,
+            wbc, lymph, age, gender, height, weight):
+
+    # Compute BMI
     try:
-        height = float(height)
-        weight = float(weight)
-        bmi = round(weight / ((height / 100) ** 2), 2) if height > 0 else "N/A"
+        bmi = round(float(weight) / ((float(height) / 100) ** 2), 2)
     except Exception:
         bmi = "N/A"
 
-    # -----------------------
+    # ----------------------------
     # Strict System Prompt
-    # -----------------------
+    # ----------------------------
     system_prompt = (
         "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"
+        "Forbidden: Any extra labs (cholesterol, vitamin D, hormones, etc.).\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"
+        "5. Tabular Mapping (Markdown table with Biomarker | Value | Range | Status | 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"
     )
@@ -111,26 +87,30 @@ def analyze(
 
     prompt = system_prompt + "\n" + patient_input
 
-    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
-    )
+    # ----------------------------
+    # Generate
+    # ----------------------------
+    try:
+        gen = pipe(
+            prompt,
+            max_new_tokens=1200,
+            temperature=0.25,
+            top_p=0.9,
+            return_full_text=False
+        )
+        generated = gen[0]["generated_text"].strip()
+    except Exception as e:
+        return f"❌ Error: {str(e)}", ""
 
-    generated = gen[0].get("generated_text", "").strip()
     if not generated:
         return "⚠️ No valid response. Please try again.", ""
 
     left_md, right_md = split_report(generated)
     return left_md, right_md
 
-# -----------------------
-# Gradio App
-# -----------------------
+# ----------------------------
+# Gradio UI
+# ----------------------------
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown("# 🏥 AI Medical Biomarker Dashboard")
 
@@ -173,4 +153,4 @@ 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)