import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
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")

# Load model directly
tokenizer = AutoTokenizer.from_pretrained("rinna/vicuna-13b-delta-finetuned-langchain-MRKL")
model = AutoModelForCausalLM.from_pretrained("rinna/vicuna-13b-delta-finetuned-langchain-MRKL")

def generate_response(input_text, chat_history=[]):
    # Tokenize the new input sentence
    new_user_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")

    # Append the new user input tokens to the chat history
    bot_input_ids = torch.cat([torch.tensor(chat_history), new_user_input_ids], dim=-1)

    # Generate a response
    chat_output = model.generate(bot_input_ids, max_length=4000, pad_token_id=tokenizer.eos_token_id)

    # Decode the response tokens into text
    response = tokenizer.decode(chat_output[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)

    return response

def chatbot_interface(input_text):
    # Generate response based on input text and chat history
    response = generate_response(input_text, chat_history)

    # Append the input and response to the chat history
    chat_history.append(tokenizer.encode(input_text + response))

    return response

chat_history = []  # Initialize chat history

iface = gr.Interface(
    fn=chatbot_interface,
    inputs=gr.inputs.Textbox(lines=2, label="Chat"),
    outputs=gr.outputs.Textbox(label="Response"),
    layout="vertical",
    title=title,
    description=description,
    examples=examples,
    theme="london"
)

iface.launch()