#!/usr/bin/python3
# -*- coding: utf-8 -*-
import argparse

import gradio as gr
import torch
from transformers import BloomTokenizerFast, BloomForCausalLM

from project_settings import project_path


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--trained_model_path',
        default=(project_path / "trained_models/bloom-1b4-sft").as_posix(),
        type=str,
    )
    parser.add_argument('--device', default='auto', type=str)

    args = parser.parse_args()
    return args


def main():
    args = get_args()

    if args.device == 'auto':
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    else:
        device = args.device

    # pretrained model
    tokenizer = BloomTokenizerFast.from_pretrained(args.trained_model_path)
    model = BloomForCausalLM.from_pretrained(args.trained_model_path)

    description = """
    FireflyBloom1b4
    
    基于 [YeungNLP/bloom-1b4-zh](https://huggingface.co/YeungNLP/bloom-1b4-zh) 预训练模型,
    基于 [YeungNLP/firefly-train-1.1M](https://huggingface.co/datasets/YeungNLP/firefly-train-1.1M) 数据集, 
    训练的等同于 [YeungNLP/firefly-bloom-1b4](https://huggingface.co/YeungNLP/firefly-bloom-1b4) 的问答模型. 
    
    训练代码是自己编写的, 在 examples 里, 总共训练了 3 个 epoch. 感觉效果还可以. 

    """

    def fn(text: str,
           max_new_tokens: int = 200,
           top_p: float = 0.85,
           temperature: float = 0.35,
           repetition_penalty: float = 1.2
           ):
        print(text)
        text = '<s>{}</s></s>'.format(text)
        input_ids = tokenizer(text, return_tensors="pt").input_ids
        input_ids = input_ids.to(device)
        outputs = model.generate(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
                                 )
        rets = tokenizer.batch_decode(outputs)
        output = rets[0].strip().replace(text, "").replace('</s>', "")
        print(output)
        return output

    demo = gr.Interface(
        fn=fn,
        inputs=[
            gr.Text(label="text"),
            gr.Number(value=200, label="max_new_tokens"),
            gr.Slider(minimum=0, maximum=1, value=0.85, label="top_p"),
            gr.Slider(minimum=0, maximum=1, value=0.35, label="temperature"),
            gr.Number(value=1.2, label="repetition_penalty"),
        ],
        outputs=[gr.Text(label="output")],
        examples=[
            [
                "将下面句子翻译成现代文：\n石中央又生一树，高百余尺，条干偃阴为五色，翠叶如盘，花径尺余，色深碧，蕊深红，异香成烟，著物霏霏。",
                200, 0.85, 0.35, 1.2
            ],
            [
                "实体识别: 1949年10月1日，人们在北京天安门广场参加开国大典。",
                200, 0.85, 0.35, 1.2
            ],
            [
                "把这句话翻译成英文: 1949年10月1日，人们在北京天安门广场参加开国大典。",
                200, 0.85, 0.35, 1.2
            ],
            [
                "晚上睡不着该怎么办. 请给点详细的介绍.",
                200, 0.85, 0.35, 1.2
            ],
            [
                "将下面的句子翻译成文言文：结婚率下降, 离婚率暴增, 生育率下降, 人民焦虑迷茫, 到底是谁的错.",
                200, 0.85, 0.35, 1.2
            ],
            [
                "对联：厌烟沿檐烟燕眼.",
                200, 0.85, 0.35, 1.2
            ],
            [
                "写一首咏雪的古诗, 标题为 \"沁园春, 雪\".",
                200, 0.85, 0.35, 1.2
            ],
        ],
        examples_per_page=50,
        title="Firefly Bloom 1b4",
        description=description,
    )
    demo.launch()

    return


if __name__ == '__main__':
    main()