ruslanmv's picture
Update app.py
83034af verified
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
import torch
import spaces
import os
IS_SPACES_ZERO = os.environ.get("SPACES_ZERO_GPU", "0") == "1"
IS_SPACE = os.environ.get("SPACE_ID", None) is not None
device = "cuda" if torch.cuda.is_available() else "cpu"
LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
print(f"Using device: {device}")
print(f"low memory: {LOW_MEMORY}")
model_name = "ruslanmv/Medical-Llama3-8B"
# Move model and tokenizer to the CUDA device
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token
@spaces.GPU
def askme(symptoms, question):
sys_message = '''\
You are an AI Medical Assistant trained on a vast dataset of health information. Please be thorough and
provide an informative answer. If you don't know the answer to a specific medical inquiry, advise seeking professional help.
'''
content = symptoms + " " + question
messages = [{"role": "system", "content": sys_message}, {"role": "user", "content": content}]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(prompt, return_tensors="pt").to(device) # Ensure inputs are on CUDA device
outputs = model.generate(**inputs, max_new_tokens=200, use_cache=True)
response_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0].strip()
# Remove system messages and content
# Extract and return the generated text, removing the prompt
# Extract only the assistant's response
# Extract only the assistant's response
answer = response_text.split('<|im_start|>assistant')[1].split('<|im_end|>')[0].strip()
return answer
# Example usage
symptoms = '''\
I'm a 35-year-old male and for the past few months, I've been experiencing fatigue,
increased sensitivity to cold, and dry, itchy skin.
'''
question = '''\
Could these symptoms be related to hypothyroidism?
If so, what steps should I take to get a proper diagnosis and discuss treatment options?
'''
examples = [
[symptoms, question]
]
css = """
/* General Container Styles */
.gradio-container {
font-family: "IBM Plex Sans", sans-serif; position: fixed; /* Ensure full-screen coverage */
top: 0;
left: 0;
width: 100vw; /* Set width to 100% viewport width */
height: 100vh; /* Set height to 100% viewport height */
margin: 0; /* Remove margins for full-screen effect */
padding: 0; /* Remove padding fol-screen background */
background-color: #212529; /* Dark background color */
color: #fff; /* Light text color for better readability */
overflow: hidden; /* Hide potential overflow content */
background-image: url("https://huggingface.co/spaces/ruslanmv/AI-Medical-Chatbot/resolve/main/notebook/local/img/background.jpg"); /* Replace with your image path */
background-size: cover; /* Stretch the image to cover the container */
background-position: center; /* Center the image horizontally and vertically */
}
/* Button Styles */
.gr-button {
color: white;
background: #007bff; /* Use a primary color for the background */
white-space: nowrap;
border: none;
padding: 10px 20px;
border-radius: 8px;
cursor: pointer;
transition: background-color 0.3s, color 0.3s;
}
.gr-button:hover {
background-color: #0056b3; /* Darken the background color on hover */
}
/* Output box styles */
.gradio-textbox {
background-color: #343a40; /* Dark background color */
color: #fff; /* Light text color for better readability */
border-color: #343a40; /* Dark border color */
border-radius: 8px;
}
"""
welcome_message = """# AI Medical Llama 3 Chatbot
Ask any medical question giving first your symptoms and get answers from our AI Medical Llama3 Chatbot
Developed by Ruslan Magana. Visit [https://ruslanmv.com/](https://ruslanmv.com/) for more information."""
symptoms_input = gr.Textbox(label="Symptoms")
question_input = gr.Textbox(label="Question")
answer_output = gr.Textbox(label="Answer")
iface = gr.Interface(
fn=askme,
inputs=[symptoms_input, question_input],
outputs=answer_output,
examples=examples,
css=css,
description=welcome_message # Add the welcome message here
)
iface.launch()
'''
with gr.Blocks(css=css) as interface:
gr.Markdown(welcome_message) # Display the welcome message
with gr.Row():
with gr.Column():
symptoms_input = gr.Textbox(label="Symptoms", placeholder="Enter symptoms here")
question_input = gr.Textbox(label="Question", placeholder="Enter question here")
generate_button = gr.Button("Ask Me", variant="primary")
with gr.Row():
answer_output = gr.Textbox(type="text", label="Answer")
interface.launch()
'''
'''
iface = gr.Interface(
fn=askme,
inputs=["text", "text"],
outputs="text",
examples=examples,
title="Medical AI Chatbot",
description="Ask me a medical question!"
)
iface.launch()
'''
'''
iface = gr.Interface(
fn=askme,
inputs=[
gr.Textbox(label="Symptoms", placeholder="Enter symptoms here"),
gr.Textbox(label="Question", placeholder="Enter question here")
],
outputs="text",
examples=examples,
title="Medical AI Chatbot",
description="Ask me a medical question!",
css=css
)
iface.launch()
'''