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

# Load model and tokenizer
model_name = "microsoft/DialoGPT-small"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

def respond(message, chat_history, chat_history_ids):
    # Encode user input
    new_input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors="pt")
    
    # Append to chat history
    if chat_history_ids is not None:
        input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
    else:
        input_ids = new_input_ids
    
    # Generate response
    chat_history_ids = model.generate(
        input_ids,
        max_length=1000,
        pad_token_id=tokenizer.eos_token_id,
        no_repeat_ngram_size=3,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.8
    )
    
    # Decode response
    response = tokenizer.decode(
        chat_history_ids[:, input_ids.shape[-1]:][0], 
        skip_special_tokens=True
    )
    
    # Update conversation history
    chat_history.append((message, response))
    
    return "", chat_history, chat_history_ids

with gr.Blocks() as demo:
    # Store model's conversation history
    state = gr.State()
    
    gr.Markdown("## DialoGPT Chatbot")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Your Message")
    clear = gr.Button("Clear History")

    msg.submit(
        respond,
        [msg, chatbot, state],
        [msg, chatbot, state]
    )
    clear.click(lambda: (None, None), outputs=[chatbot, state], queue=False)

demo.launch()