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


title = "????AI ChatBot"
description = "A State-of-the-Art Large-scale Pretrained Response generation model (DialoGPT)"
examples = [["How are you?"]]


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


def predict(user_input, history=[]):
    # Tokenize user input + end of string
    new_user_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")

    # Concatenate with history (convert to tensor if not empty)
    bot_input_ids = torch.cat([torch.tensor(history, dtype=torch.long), new_user_input_ids], dim=-1) if history else new_user_input_ids

    # Generate response
    output_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    history = output_ids.tolist()

    # Decode only new tokens (i.e., skip input tokens)
    response = tokenizer.decode(output_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Return chatbot-friendly tuple and updated history
    return [(user_input, response)], history

gr.Interface(
    fn=predict,
    inputs=[gr.Textbox(placeholder="Say something..."), gr.State()],
    outputs=[gr.Chatbot(), gr.State()],
    title="🤖 AI ChatBot",
    description="A State-of-the-Art Large-scale Pretrained Response Generation Model (DialoGPT)",
    examples=[["How are you?"]],
    theme="finlaymacklon/boxy_violet",
).launch()