import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")
model = AutoModelForSeq2SeqLM.from_pretrained("ahmed792002/Finetuning_T5_HealthCare_Chatbot")

# Define the chatbot function
def chatbot(input_text):
    inputs = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True)
    outputs = model.generate(inputs["input_ids"], max_length=100, num_beams=4, early_stopping=True)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Set up the Gradio interface
interface = gr.Interface(
    fn=chatbot,
    inputs=gr.inputs.Textbox(label="Enter your query"),
    outputs=gr.outputs.Textbox(label="Response"),
    title="Healthcare Chatbot",
    description="Ask healthcare-related questions, and get responses from the fine-tuned T5 model."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()