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

tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium")

chat_history_ids = None

def chat(input_text, history=[]):
    global chat_history_ids
    new_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors='pt')

    if chat_history_ids is not None:
        bot_input_ids = torch.cat([chat_history_ids, new_input_ids], dim=-1)
    else:
        bot_input_ids = new_input_ids

    chat_history_ids = model.generate(
        bot_input_ids,
        max_length=1000,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_k=50,
        top_p=0.95,
        temperature=0.75,
    )

    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    history.append((input_text, response))
    return history, history

with gr.Blocks() as demo:
    gr.Markdown("## 💖 AI Girlfriend Chat (DialoGPT)")
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Type something...")
    state = gr.State([])

    msg.submit(chat, [msg, state], [chatbot, state])

demo.launch()