import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load the model and tokenizer
model_name = "google/flan-t5-xxl"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Define the chat function
def chat(message):
    # Encode the user's message
    inputs = tokenizer.encode(message, return_tensors="pt")
    # Generate a response from the model
    outputs = model.generate(inputs, max_length=1024, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Return the response
    return response

# Set up the Gradio interface
block = gr.Blocks(css=".gradio-container {background-color: lightgray}")

with block:
    with gr.Row():
        gr.Markdown("<h3><center>SplitticAI Chatbot</center></h3>")

    chatbot = gr.Chatbot()

    with gr.Row():
        message = gr.Textbox(
            label="What's your question?",
            placeholder="What would you like to ask me?",
            lines=1,
        )
        submit = gr.Button(value="Send", variant="secondary").style(full_width=False)

    gr.Examples(
        examples=[
            "What is artificial intelligence?",
            "How does SplitticAI work?",
            "Can you tell me a joke?",
        ],
        inputs=message,
    )

    gr.HTML("Ask SplitticAI anything and get an answer!")

    gr.HTML(
        "<center>Powered by <a href='https://huggingface.co/google/flax-t5-xxl-qa-121k'>google/flax-t5-xxl-qa-121k</a></center>"
    )

    state = gr.State()
    agent_state = gr.State()

    submit.click(chat, inputs=[message], outputs=[chatbot])

block.launch(debug=True)