app.py

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"
#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
#dtype = torch.float16
LOW_MEMORY = os.getenv("LOW_MEMORY", "0") == "1"
print(f"Using device: {device}")
#print(f"Using dtype: {dtype}")
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() #skip_special_tokens=True
    # Remove system messages and content
   #  Extract only the assistant's response
    assistant_response =response_text.split("assistant")[1].strip().split("user")[0].strip()   
    return assistant_response



# 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]
]

iface = gr.Interface(
    fn=askme,
    inputs=["text", "text"],
    outputs="text",
    examples=examples,
    title="Medical AI Chatbot",
    description="Ask me a medical question!"
)

iface.launch()
'''


# Read CSS from file
css = """
#content_align {
    text-align: center;
    padding: 20px;
    background-color: #282c34;
    color: #fff;
    font-family: Arial, sans-serif;
}
"""

# The welcome message with improved styling
welcome_message = '''
<div id="content_align">
  <h1>AI Medical Chatbot</h1>
  <p>Ask any medical question and get answers from our AI Medical Chatbot</p>
  <p>Developed by Ruslan Magana. Visit <a href="https://ruslanmv.com/">https://ruslanmv.com/</a> for more information.</p>
</div>
'''

# Creating Gradio interface with full-screen styling
iface = gr.Interface(
    fn=askme,
    inputs=["text", "text"],
    outputs="text",
    title="Medical AI Chatbot",
    description="Ask me a medical question!",
    layout="vertical",
    css=css
)
# Display the welcome message
iface.add_text(welcome_message)
# Input elements
symptoms_input = iface.add_textbox("text", label="Symptoms", placeholder="Enter symptoms here")
question_input = iface.add_textbox("text", label="Question", placeholder="Enter question here")
# Output element
answer_output = iface.add_textbox("text", label="Answer")

# Assuming you have a function `askme` that processes the symptoms and question and returns a response
iface.launch(examples=examples)