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

# Load models and tokenizers
def load_models():
    # Load a conversational model and tokenizer (you can customize it further)
    model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")
    tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
    return model, tokenizer

# Generate responses
def chat_with_model(user_input, model, tokenizer, chat_history):
    # Tokenize the user input and chat history
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')

    # Append new user input to chat history
    bot_input_ids = torch.cat([chat_history, new_user_input_ids], dim=-1) if chat_history is not None else new_user_input_ids

    # Generate a response from the model
    chat_history = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)

    # Decode the model's output and return
    bot_output = tokenizer.decode(chat_history[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    return chat_history, bot_output

# Initialize model and tokenizer
model, tokenizer = load_models()

# Build Gradio interface
def build_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Chagrin AI Chatbot")

        # Set up chat window
        chatbot = gr.Chatbot()

        # Create text input box for user to type
        user_input = gr.Textbox(label="Type your message", placeholder="Ask something...", interactive=True)

        # Create button for sending the input
        submit_btn = gr.Button("Send Message")

        # Button click function
        submit_btn.click(chat_with_model, inputs=[user_input, model, tokenizer, chatbot], outputs=[chatbot, chatbot])

    demo.launch()

# Run the Gradio interface
if __name__ == "__main__":
    build_gradio_interface()