import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "physician-ai/llama3-8b-finetuned"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name,device_map="auto")

# System prompt to guide the chatbot's responses
system_prompt = "You are a helpful medical assistant. Answer the following questions based on medical knowledge."

def chat(user_input):
    # Combine system prompt with user input
    prompt = system_prompt + "\nUser: " + user_input + "\nBot:"
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(inputs, max_length=200, num_return_sequences=1)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("Bot:")[-1].strip()

iface = gr.Interface(fn=chat, inputs="text", outputs="text", title="Medical Chatbot")
iface.launch()