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

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

# Track chat history across calls
chat_history_ids = None

def chatbot(user_input):
    global chat_history_ids
    
    # Encode user input + eos
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors='pt')
    
    # Append new user input to chat history
    if chat_history_ids is not None:
        bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1)
    else:
        bot_input_ids = new_user_input_ids
    
    # Generate response adding ~50 tokens
    chat_history_ids = model.generate(
        bot_input_ids,
        max_length=bot_input_ids.shape[-1]+50,
        pad_token_id=tokenizer.eos_token_id,
        do_sample=True,
        top_k=50,
        temperature=0.7
    )
    
    # Decode only new tokens
    response = tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response

iface = gr.Interface(
    fn=chatbot,
    inputs="text",
    outputs="text",
    title="DialoGPT Chatbot"
)
iface.launch()