Update app.py

app.py

@@ -1,95 +1,113 @@
-import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# Load your model from Hugging Face Hub
-MODEL_ID = "Muhammadidrees/MedicalInsights"
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto")
-pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-
-# Function to build structured input and query the LLM
-def analyze(
-    albumin, creatinine, glucose, crp, mcv, rdw, alp,
-    wbc, lymph, age, gender, height, weight, bmi
-):
-    # System-style instruction
-    system_prompt = (
-        "You are an advanced AI medical assistant. "
-        "Analyze the patient’s biomarkers and demographics. "
-        "Provide a structured assessment including: "
-        "patient_profile, lab_results, risk_assessment, clinical_impression, recommendations. "
-    )
-
-    # 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} %
-    """
-
-    prompt = system_prompt + "\n" + patient_input
-
-    # Call LLM
-    result = pipe(prompt, max_new_tokens=400, do_sample=True, temperature=0.6)
-    return result[0]["generated_text"]
-
-
-# Build Gradio UI
-with gr.Blocks() as demo:
-    gr.Markdown("## 🧪 Medical Insights AI — Enter Patient Data")
-
-    with gr.Row():
-        albumin = gr.Number(label="Albumin (g/dL)")
-        wbc = gr.Number(label="White Blood Cell Count (K/uL)")
-
-    with gr.Row():
-        creatinine = gr.Number(label="Creatinine (mg/dL)")
-        lymph = gr.Number(label="Lymphocyte Percentage (%)")
-
-    with gr.Row():
-        glucose = gr.Number(label="Glucose (mg/dL)")
-        age = gr.Number(label="Age (years)")
-
-    with gr.Row():
-        crp = gr.Number(label="C-Reactive Protein (mg/L)")
-        gender = gr.Dropdown(choices=["Male", "Female"], label="Gender")
-
-    with gr.Row():
-        mcv = gr.Number(label="Mean Cell Volume (fL)")
-        height = gr.Number(label="Height (cm)")
-
-    with gr.Row():
-        rdw = gr.Number(label="Red Cell Distribution Width (%)")
-        weight = gr.Number(label="Weight (kg)")
-
-    with gr.Row():
-        alp = gr.Number(label="Alkaline Phosphatase (U/L)")
-        bmi = gr.Number(label="BMI")
-
-    analyze_btn = gr.Button("🔎 Analyze")
-    output = gr.Textbox(label="AI Medical Assessment", lines=12)
-
-    analyze_btn.click(
-        fn=analyze,
-        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
-                wbc, lymph, age, gender, height, weight, bmi],
-        outputs=output
-    )
-
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+
+MODEL_ID = "Muhammadidrees/MedicalInsights"
+
+# Create offload directory inside Space
+offload_dir = "./offload"
+os.makedirs(offload_dir, exist_ok=True)
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",          # still balances CPU/GPU
+    offload_folder=offload_dir, # where to dump layers too big for RAM
+    trust_remote_code=True
+)
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device_map="auto"
+)
+
+# Function to build structured input and query the LLM
+def analyze(
+    albumin, creatinine, glucose, crp, mcv, rdw, alp,
+    wbc, lymph, age, gender, height, weight, bmi
+):
+    # System-style instruction
+    system_prompt = (
+        "You are an advanced AI medical assistant. "
+        "Analyze the patient’s biomarkers and demographics. "
+        "Provide a structured assessment including: "
+        "patient_profile, lab_results, risk_assessment, clinical_impression, recommendations. "
+    )
+
+    # 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} %
+    """
+
+    prompt = system_prompt + "\n" + patient_input
+
+    # Call LLM
+    result = pipe(prompt, max_new_tokens=400, do_sample=True, temperature=0.6)
+    return result[0]["generated_text"]
+
+
+# Build Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧪 Medical Insights AI — Enter Patient Data")
+
+    with gr.Row():
+        albumin = gr.Number(label="Albumin (g/dL)")
+        wbc = gr.Number(label="White Blood Cell Count (K/uL)")
+
+    with gr.Row():
+        creatinine = gr.Number(label="Creatinine (mg/dL)")
+        lymph = gr.Number(label="Lymphocyte Percentage (%)")
+
+    with gr.Row():
+        glucose = gr.Number(label="Glucose (mg/dL)")
+        age = gr.Number(label="Age (years)")
+
+    with gr.Row():
+        crp = gr.Number(label="C-Reactive Protein (mg/L)")
+        gender = gr.Dropdown(choices=["Male", "Female"], label="Gender")
+
+    with gr.Row():
+        mcv = gr.Number(label="Mean Cell Volume (fL)")
+        height = gr.Number(label="Height (cm)")
+
+    with gr.Row():
+        rdw = gr.Number(label="Red Cell Distribution Width (%)")
+        weight = gr.Number(label="Weight (kg)")
+
+    with gr.Row():
+        alp = gr.Number(label="Alkaline Phosphatase (U/L)")
+        bmi = gr.Number(label="BMI")
+
+    analyze_btn = gr.Button("🔎 Analyze")
+    output = gr.Textbox(label="AI Medical Assessment", lines=12)
+
+    analyze_btn.click(
+        fn=analyze,
+        inputs=[albumin, creatinine, glucose, crp, mcv, rdw, alp,
+                wbc, lymph, age, gender, height, weight, bmi],
+        outputs=output
+    )
+
 demo.launch()