ruslanmv's picture
Update app.py
83034af verified
raw
history blame contribute delete
No virus
5.5 kB
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()
'''