from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM

#tokenizer = AutoTokenizer.from_pretrained("bigscience/T0pp")
#model = AutoModelForSeq2SeqLM.from_pretrained("bigscience/T0pp")
tokenizer = AutoTokenizer.from_pretrained("ethzanalytics/ai-msgbot-gpt2-M")
model = AutoModelForCausalLM.from_pretrained("ethzanalytics/ai-msgbot-gpt2-M")

import gradio as gr
import random

def chat(message):
    history = gr.get_state() or []
    if message.startswith("How many"):
        response = random.randint(1,10)
    elif message.startswith("How"):
        response = random.choice(["Great", "Good", "Okay", "Bad"])
    elif message.startswith("Where"):
        response = random.choice(["Here", "There", "Somewhere"])
    else:
        inputs = tokenizer.encode(message, return_tensors="pt")
        input_len = len(message)
        outputs = model.generate(inputs)
        response = tokenizer.decode(outputs[0])[input_len:]

    history.append((message, response))
    gr.set_state(history)
    html = "<div class='chatbot'>"
    for user_msg, resp_msg in history:
        html += f"<div class='user_msg'>{user_msg}</div>"
        html += f"<div class='resp_msg'>{resp_msg}</div>"
    html += "</div>"
    return html

iface = gr.Interface(chat, "text", "html", css="""
    .chatbox {display:flex;flex-direction:column}
    .user_msg, .resp_msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}
    .user_msg {background-color:cornflowerblue;color:white;align-self:start}
    .resp_msg {background-color:lightgray;align-self:self-end}
""", allow_screenshot=False, allow_flagging=False)
iface.launch()