import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model and tokenizer
MODEL_NAME = "microsoft/phi-4"
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

def chatbot_response(user_input, chat_history=[]):
    """Generates a response from the chatbot model."""
    # Tokenize input and add chat history
    input_ids = tokenizer.encode(user_input, return_tensors="pt")

    # Generate response
    output = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response = tokenizer.decode(output[:, input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Update chat history
    chat_history.append((user_input, response))
    return chat_history, "\n".join([f"You: {msg}\nBot: {res}" for msg, res in chat_history])

# Gradio Interface
with gr.Blocks() as chatbot_ui:
    gr.Markdown("## Chatbot Interface")

    chat_history = gr.State([])  # Stores the chat history

    with gr.Row():
        user_input = gr.Textbox(placeholder="Type your message here...", label="Your Input")
        submit_button = gr.Button("Send")

    with gr.Row():
        chat_display = gr.Textbox(label="Chat History", lines=20, interactive=False)

    # Event listener
    submit_button.click(chatbot_response, inputs=[user_input, chat_history], outputs=[chat_history, chat_display])

chatbot_ui.launch()