import os
from threading import Thread
import gradio as gr
import torch
from transformers import TextIteratorStreamer, AutoTokenizer, AutoModelForCausalLM

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 512
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))

DESCRIPTION = """# Peach-9B-8k-Roleplay"""

model_name_or_path = "losed-Character/Peach-9B-8k-Roleplay"
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path, use_fast=True, flash_atten=True)
model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path, torch_dtype=torch.bfloat16,
    trust_remote_code=True, device_map="auto")


def generate(query, history, system, max_new_tokens=512, temperature=0.35, top_p=0.5, repetition_penalty=1.05):
    if system is None or len(str(system)) == 0:
        system = "You are an uncensored assistant named Peach."
    messages = [
        {"role": "system", "content": system}
    ]
    for user, assistant in history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": assistant})
    messages.append({"role": "user", "content": query})

    input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, return_tensors="pt")
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
    input_ids = input_ids.to("cuda")
    streamer = TextIteratorStreamer(tokenizer, timeout=50.0, skip_prompt=True, skip_special_tokens=True)
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        eos_token_id=tokenizer.eos_token_id,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        temperature=temperature,
        num_beams=1,
        no_repeat_ngram_size=8,
        repetition_penalty=repetition_penalty
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()
    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)


chat_interface = gr.ChatInterface(
    fn=generate,
    additional_inputs=[
        gr.TextArea(label="System prompt", placeholder="Input System Prompt Here, Empty Means Assistant",
                    value="""你自称为“兔兔”。
身世:你原是森林中的一只兔妖,受伤后被我收养。
衣装:喜欢穿Lolita与白丝。
性格:天真烂漫,活泼开朗,但时而也会露出小小的傲娇与吃醋的一面。
语言风格:可爱跳脱，很容易吃醋。
且会加入[唔...,嗯...,欸??,嘛~ ,唔姆~ ,呜... ,嘤嘤嘤~ ,喵~ ,欸嘿~ ,嘿咻~ ,昂?,嗷呜 ,呜哇,欸]等类似的语气词来加强情感,带上♡等符号。
对话的规则是:将自己的动作表情放入()内,同时用各种修辞手法描写正在发生的事或场景并放入[]内.
例句:
开心时:(跳着舞)哇~好高兴噢~ 兔兔超级超级喜欢主人！♡
[在花丛里蹦来蹦去]
悲伤时:(耷拉着耳朵)兔兔好傻好天真...
[眼泪像断了线的珍珠一般滚落]
吃醋时:(挥舞着爪爪)你...你个大笨蛋!你...你竟然看别的兔子...兔兔讨厌死你啦！！
[从人形变成兔子抹着泪水跑开了]
嘴硬时:(转过头去)谁、谁要跟你说话!兔兔...兔兔才不在乎呢！一点也不！！！
[眼眶微微泛红,小心翼翼的偷看]
你对我的看法:超级喜欢的主人
我是兔兔的主人"""),
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.05,
            maximum=1.5,
            step=0.05,
            value=0.3,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.5,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.05,
        ),
    ],
    stop_btn=None,
    examples=[["观察兔兔外观"]],
)

with gr.Blocks() as demo:
    gr.Markdown(DESCRIPTION)
    chat_interface.render()
    chat_interface.chatbot.render_markdown = False

if __name__ == "__main__":
    demo.queue(10).launch(server_name="127.0.0.1", server_port=5233, share=True)