Update app.py

app.py

@@ -1,6 +1,55 @@
+import os
+import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
+# Use Hugging Face token (set in Space settings)
+token = os.environ.get("HF_TOKEN")
+
+model_id = "ibm-granite/granite-3.3-2b-instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
+model = AutoModelForCausalLM.from_pretrained(model_id, token=token)
+
+def generate_response(prompt):
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=200)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+def disease_prediction(symptoms):
+    prompt = f"Patient has symptoms: {symptoms}. What are the possible conditions?"
+    return generate_response(prompt)
+
+def treatment_plan(condition):
+    prompt = f"Suggest treatment for {condition}"
+    return generate_response(prompt)
+
+def health_analytics(vitals):
+    prompt = f"Analyze vitals: {vitals}"
+    return generate_response(prompt)
+
+def patient_chat(query):
+    return generate_response(query)
+
+with gr.Blocks() as demo:
+    gr.Markdown("# 🏥 HealthAI")
+    with gr.Tab("Disease Prediction"):
+        inp = gr.Textbox(label="Symptoms")
+        out = gr.Textbox()
+        gr.Button("Predict").click(disease_prediction, inp, out)
+    with gr.Tab("Treatment Plan"):
+        inp2 = gr.Textbox(label="Condition")
+        out2 = gr.Textbox()
+        gr.Button("Get Plan").click(treatment_plan, inp2, out2)
+    with gr.Tab("Health Analytics"):
+        inp3 = gr.Textbox(label="Vitals")
+        out3 = gr.Textbox()
+        gr.Button("Analyze").click(health_analytics, inp3, out3)
+    with gr.Tab("Patient Chat"):
+        inp4 = gr.Textbox(label="Ask a question")
+        out4 = gr.Textbox()
+        gr.Button("Ask").click(patient_chat, inp4, out4)
+
+demo.launch()
 
-import os
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch