HealthAi / app.py
reshma-05's picture
Update app.py
d4d4a40 verified
import os
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# βœ… IBM Granite model setup
model_id = "ibm-granite/granite-3.3-2b-instruct"
token = os.getenv("HF_TOKEN") # Ensure your Hugging Face token is set in the environment
# βœ… Load model and tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id, token=token)
model = AutoModelForCausalLM.from_pretrained(
model_id,
token=token,
device_map="auto",
torch_dtype=torch.float32
)
# βœ… Query function
def query_granite(prompt):
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=100)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
# βœ… Gradio UI using Tabs (instead of Pages)
with gr.Blocks() as demo:
gr.Markdown("# πŸ₯ Welcome to HealthAI")
gr.Markdown("Your intelligent healthcare assistant.")
with gr.Tab("🩺 Symptoms"):
def identify(symptom):
return query_granite(f"What illness could cause: {symptom}?")
symptom = gr.Textbox(label="Enter your symptom")
output = gr.Textbox(label="AI Diagnosis")
btn = gr.Button("Analyze")
btn.click(identify, inputs=symptom, outputs=output)
with gr.Tab("🌿 Remedies"):
def get_remedies(issue):
return query_granite(f"What are home remedies for {issue}?")
issue = gr.Textbox(label="What are you suffering from?")
remedy_output = gr.Textbox(label="Suggested Remedy")
remedy_btn = gr.Button("Suggest")
remedy_btn.click(get_remedies, inputs=issue, outputs=remedy_output)
with gr.Tab("πŸ₯— Diet"):
def suggest(goal):
return query_granite(f"Suggest a diet for: {goal}")
goal = gr.Textbox(label="Your health goal")
diet_output = gr.Textbox(label="Diet Plan")
diet_btn = gr.Button("Get Plan")
diet_btn.click(suggest, inputs=goal, outputs=diet_output)
with gr.Tab("🧠 Mental Wellness"):
def tip(topic):
return query_granite(f"Mental health advice about: {topic}")
topic = gr.Textbox(label="Enter mental health topic")
tip_output = gr.Textbox(label="Wellness Tip")
tip_btn = gr.Button("Get Tip")
tip_btn.click(tip, inputs=topic, outputs=tip_output)
with gr.Tab("❓ FAQs"):
gr.Markdown("### ❓ FAQs")
gr.Markdown("**Q1:** What is HealthAI?")
gr.Markdown("**A:** It's an AI assistant to help with health-related queries using IBM Granite 3.3-2B.")
demo.launch()