"""
Try out gradio.Chatinterface.

colab gradio-chatinterface.

%%writefile reuirements.txt
gradio
transformers
sentencepiece
torch

"""
# pylint: disable=line-too-long, missing-module-docstring, missing-function-docstring
# import torch
import gradio as gr
from examples_list import examples_list
from transformers import AutoModel, AutoTokenizer  # AutoModelForCausalLM,

# device = "cuda" if torch.cuda.is_available() else "cpu"

# tokenizer = AutoTokenizer.from_pretrained("stabilityai/StableBeluga2", use_fast=False)
# model = AutoModelForCausalLM.from_pretrained("stabilityai/StableBeluga2", torch_dtype=torch.float16, low_cpu_mem_usage=True, device_map="auto")
# system_prompt = "### System:\nYou are Stable Beluga, an AI that follows instructions extremely well. Help as much as you can. Remember, be safe, and don't do anything illegal.\n\n"
# pipeline = pipeline(task="text-generation", model="meta-llama/Llama-2-7b")
tokenizer = AutoTokenizer.from_pretrained(
    "THUDM/chatglm2-6b-int4", trust_remote_code=True
)
chat_model = AutoModel.from_pretrained(
    "THUDM/chatglm2-6b-int4", trust_remote_code=True  # 3.92G
).float()


def chat(message, history):
    # prompt = f"{system_prompt}### User: {message}\n\n### Assistant:\n"
    # inputs = tokenizer(prompt, return_tensors="pt").to(device=device)
    # output = model.generate(**inputs, do_sample=True, top_p=0.95, top_k=0, max_new_tokens=256)
    # return tokenizer.decode(output[0], skip_special_tokens=True)
    for response, _ in chat_model.stream_chat(
        tokenizer, message, history, max_length=2048, top_p=0.7, temperature=0.95
    ):
        yield response

chatbot = gr.Chatbot([], label="Bot", height=450)
textbox = gr.Textbox('', scale=10, label='', lines=2, placeholder="Ask me anything")
submit_btn = gr.Button(value="Send", scale=1, min_width=0, variant="primary")

interf = gr.ChatInterface(
    chat,
    chatbot=chatbot,
    textbox=textbox,
    submit_btn=submit_btn,
    title="gradio-chatinterface-tryout",
    examples=examples_list,
    theme=gr.themes.Glass(text_size="sm", spacing_size="sm"),
).queue(max_size=5)


if __name__ == "__main__":
    interf.launch(debug=True)