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

def flatten(l):
    return [item for sublist in l for item in sublist]

tokenizer = AutoTokenizer.from_pretrained("model")
model = AutoModelWithLMHead.from_pretrained("model")

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox()
    clear = gr.ClearButton([msg, chatbot])

    def respond(message, chat_history):
        input_ids = tokenizer.encode(message + tokenizer.eos_token, return_tensors='pt')
        if len(chat_history):
            tokenized_chat_history = [tokenizer.encode(x + tokenizer.eos_token, return_tensors='pt')[0] for x in flatten(chat_history)]
            tokenized_chat_history = torch.cat(tokenized_chat_history).unsqueeze(0)
        bot_input_ids = torch.cat([tokenized_chat_history, input_ids], dim=-1) if len(chat_history) else input_ids
        output = model.generate(
            bot_input_ids,
            max_new_tokens=50, 
            pad_token_id=tokenizer.eos_token_id,
            do_sample=True,
            top_k=50,
            top_p=0.95,
        )
        bot_message = str(tokenizer.decode(output[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True))
        chat_history.append((message, bot_message))
        return "", chat_history
    msg.submit(respond, [msg, chatbot], [msg, chatbot])

demo.launch()