import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline

model_id = "captainali01/pediatric-chatbot"

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)

pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

def respond(message, history=[]):
    response = pipe(
        message,
        max_new_tokens=256,
        temperature=0.7,
        top_p=0.95,
        do_sample=True
    )[0]["generated_text"]
    return response

chatbot = gr.ChatInterface(
    fn=respond,
    title="Pediatric Chatbot",
    description="A fine-tuned pediatric assistant powered by your custom model.",
)

if __name__ == "__main__":
    chatbot.launch()
#