Update app.py

app.py

@@ -1,30 +1,52 @@
 import gradio as gr
 from transformers import pipeline
 
-# Load the model from Hugging Face Hub
-generator = pipeline("text-generation", model="ibm-granite/granite-3.3-2b-instruct", device_map="auto")
+# Load the model with accelerate support
+generator = pipeline(
+    "text-generation", 
+    model="ibm-granite/granite-3.3-2b-instruct", 
+    device_map="auto"
+)
 
+# Predict disease from symptoms
 def predict_disease(symptoms):
     prompt = f"User has the following symptoms: {symptoms}. What is the most probable disease?"
-    result = generator(prompt, max_new_tokens=100, do_sample=True)[0]["generated_text"]
-    return result.split(prompt)[-1].strip()
+    try:
+        output = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
+        if prompt in output:
+            response = output.split(prompt)[-1].strip()
+        else:
+            response = output.strip()
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
 
+# Suggest home remedy from disease
 def home_remedy(disease):
     prompt = f"Suggest a natural home remedy for the disease: {disease}."
-    result = generator(prompt, max_new_tokens=100, do_sample=True)[0]["generated_text"]
-    return result.split(prompt)[-1].strip()
+    try:
+        output = generator(prompt, max_new_tokens=150, do_sample=True, temperature=0.7)[0]["generated_text"]
+        if prompt in output:
+            response = output.split(prompt)[-1].strip()
+        else:
+            response = output.strip()
+        return response
+    except Exception as e:
+        return f"Error: {str(e)}"
 
+# UI with Gradio
 with gr.Blocks() as demo:
-    gr.Markdown("## 🧠 HealthAI")
+    gr.Markdown("# 🧠 HealthAI")
+    gr.Markdown("Identify diseases from symptoms and get natural home remedies.")
 
     with gr.Tab("🔍 Symptoms Identifier"):
-        symptoms_input = gr.Textbox(label="Enter your symptoms")
+        symptoms_input = gr.Textbox(label="Enter your symptoms", placeholder="e.g. fever, cough, sore throat")
         symptoms_output = gr.Textbox(label="Predicted Disease")
         symptoms_btn = gr.Button("Predict Disease")
         symptoms_btn.click(fn=predict_disease, inputs=symptoms_input, outputs=symptoms_output)
 
     with gr.Tab("🌿 Home Remedies"):
-        disease_input = gr.Textbox(label="Enter your disease")
+        disease_input = gr.Textbox(label="Enter your disease", placeholder="e.g. Cold, Headache")
         remedy_output = gr.Textbox(label="Suggested Remedy")
         remedy_btn = gr.Button("Get Remedy")
         remedy_btn.click(fn=home_remedy, inputs=disease_input, outputs=remedy_output)