#!/usr/bin/python3
# -*- coding: utf-8 -*-
from typing import List, Tuple
from threading import Thread

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

from project_settings import project_path


def greet(question: str, history: List[Tuple[str, str]]):
    answer = "Hello " + question + "!"
    result = history + [(question, answer)]
    return result


def chat_with_llm_non_stream(question: str,
                             history: List[Tuple[str, str]],
                             pretrained_model_name_or_path: str,
                             max_new_tokens: int, top_p: float, temperature: float, repetition_penalty: float,
                             ):
    device = "cuda" if torch.cuda.is_available() else "cpu"

    model = AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        offload_folder="./offload",
        offload_state_dict=True,
        # load_in_4bit=True,
    )
    model = model.to(device)
    model = model.bfloat16().eval()

    tokenizer = AutoTokenizer.from_pretrained(
        pretrained_model_name_or_path,
        trust_remote_code=True,
        # llama不支持fast
        use_fast=False if model.config.model_type == "llama" else True,
        padding_side="left"
    )

    # QWenTokenizer比较特殊, pad_token_id, bos_token_id, eos_token_id 均 为None. eod_id对应的token为<|endoftext|>
    if tokenizer.__class__.__name__ == "QWenTokenizer":
        tokenizer.pad_token_id = tokenizer.eod_id
        tokenizer.bos_token_id = tokenizer.eod_id
        tokenizer.eos_token_id = tokenizer.eod_id

    input_ids = tokenizer(
        question,
        return_tensors="pt",
        add_special_tokens=False,
    ).input_ids.to(device)
    bos_token_id = torch.tensor([[tokenizer.bos_token_id]], dtype=torch.long).to(device)
    eos_token_id = torch.tensor([[tokenizer.eos_token_id]], dtype=torch.long).to(device)
    input_ids = torch.concat([bos_token_id, input_ids, eos_token_id], dim=1)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=True,
            top_p=top_p,
            temperature=temperature,
            repetition_penalty=repetition_penalty,
            eos_token_id=tokenizer.eos_token_id
        )
        outputs = outputs.tolist()[0][len(input_ids[0]):]
        response = tokenizer.decode(outputs)
        response = response.strip().replace(tokenizer.eos_token, "").strip()

    return


def main():
    description = """
    chat llm
    """

    with gr.Blocks() as blocks:
        gr.Markdown(value="gradio demo")

        chatbot = gr.Chatbot([], elem_id="chatbot", height=400)
        with gr.Row():
            with gr.Column(scale=4):
                text_box = gr.Textbox(show_label=False, placeholder="Enter text and press enter", container=False)
            with gr.Column(scale=1):
                submit_button = gr.Button("💬Submit")
            with gr.Column(scale=1):
                clear_button = gr.Button(
                    '🗑️Clear',
                    variant='secondary',
                )

        with gr.Row():
            with gr.Column(scale=1):
                max_new_tokens = gr.Slider(minimum=0, maximum=512, value=512, step=1, label="max_new_tokens"),
            with gr.Column(scale=1):
                top_p = gr.Slider(minimum=0, maximum=1, value=0.85, step=0.01, label="top_p"),
            with gr.Column(scale=1):
                temperature = gr.Slider(minimum=0, maximum=1, value=0.35, step=0.01, label="temperature"),
            with gr.Column(scale=1):
                repetition_penalty = gr.Slider(minimum=0, maximum=2, value=1.2, step=0.01, label="repetition_penalty"),

        with gr.Row():
            model_name = gr.Dropdown(choices=["Qwen/Qwen-7B-Chat"],
                                     value="Qwen/Qwen-7B-Chat",
                                     label="model_name",
                                     )
        gr.Examples(examples=["你好"], inputs=text_box)

        inputs = [
            text_box, chatbot, model_name,
            max_new_tokens, top_p, temperature, repetition_penalty
        ]
        outputs = [
            chatbot
        ]
        text_box.submit(chat_with_llm_non_stream, inputs, outputs)
        submit_button.click(chat_with_llm_non_stream, inputs, outputs)
        clear_button.click(
            fn=lambda: ('', ''),
            outputs=[text_box, chatbot],
            queue=False,
            api_name=False,
        )

    blocks.queue().launch()

    return


if __name__ == '__main__':
    main()