import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

model_name = "google/flan-t5-base"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def chat(user_input, history):

    prompt = f"""
    You are a helpful assistant.
    Answer clearly and correctly.

    User question: {user_input}
    """

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)

    outputs = model.generate(
        **inputs,
        max_length=200,
        temperature=0.7
    )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    history.append((user_input, response))

    return history, history


demo = gr.Interface(
    fn=chat,
    inputs=[
        gr.Textbox(label="Ask something"),
        gr.State([])
    ],
    outputs=[
        gr.Chatbot(),
        gr.State()
    ],
    title="Correct AI Chatbot 🤖 (Fixed Model)"
)

demo.launch()