from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers import AutoTokenizer, pipeline

import gradio as gr
import os



# Initialize message history array
message_history = []
initial_message = "Please write your prompt here and press 'enter'"
device = "cuda" # the device to load the model onto

model = AutoModelForCausalLM.from_pretrained(
    "DuongTrongChi/Rikka-1.8B-v2.2",
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained("DuongTrongChi/Rikka-1.8B-v2.2")
eos_token = tokenizer("<|im_end|>",add_special_tokens=False)["input_ids"][0]
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)


def chat(prompt):
    def chat_template(prompt):
        system_prompt = """Bạn là một trợ lý hữu ích, tôn trọng và trung thực. Luôn trả lời một cách hữu ích nhất có thể trong khi vẫn an toàn. Câu trả lời của bạn không được bao gồm bất kỳ nội dung có hại, phi đạo đức, phân biệt chủng tộc, phân biệt giới tính, độc hại, nguy hiểm hoặc bất hợp pháp. """
        return [
            {"role": "system", "content": system_prompt},
            {"role": "question", "content": prompt}
        ]

    prompt = pipe.tokenizer.apply_chat_template(chat_template(prompt), tokenize=False, add_generation_prompt=True)
    outputs = pipe(prompt, max_new_tokens=512, do_sample=True, temperature=0.3, top_k=50, top_p=0.95, eos_token_id=eos_token, pad_token_id=eos_token)
    return outputs[0]['generated_text'][len(prompt):].strip()



with gr.Blocks(theme='abidlabs/dracula_test') as chatblock:

    gr.Markdown("<h1><center>Welcome to my personal AI assistant (powered by OpenAI)</center></h1>")

    Chatbot = gr.Chatbot()
    with gr.Row():
        txt = gr.Textbox(
            show_label=False,
            placeholder = initial_message).style(container=False)
        state = gr.State()
        txt.submit(chat, txt, Chatbot)
        txt.submit(None, None, txt, _js="() => {''}")

chatblock.launch()