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

# Load pre-trained model and tokenizer
model_name = "microsoft/DialoGPT-large" # use small, medium
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Store chat history
chat_history = []

def chat_with_bot(user_input):
    global chat_history
    # Tokenize input with chat history
    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")

    # Append previous chat history if available
    if chat_history:
        input_ids = torch.cat([chat_history[-1], input_ids], dim=-1)

    # Generate a response
    response_ids = model.generate(input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)
    response_text = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)

    # Save history
    chat_history.append(input_ids)

    return response_text

# Create Gradio Interface
app = gr.Interface(fn=chat_with_bot, inputs="text", outputs="text", title="AI Chatbot")

# Launch the app
app.launch()