# coding=utf-8
# author: xusong <xusong28@jd.com>
# time: 2022/8/23 12:58

import gradio as gr

import torch
import modeling_kplug_s2s_patch
from transformers import BertTokenizer, BartModel, BartForConditionalGeneration



# 改成 huggingface-model自动模型
model_dir = "models/ft_cepsum_jiadian/"
model = BartForConditionalGeneration.from_pretrained(model_dir)  # cnn指的是cnn daily mail
tokenizer = BertTokenizer.from_pretrained(model_dir)


def generate(text):
    inputs = tokenizer([text], max_length=512, return_tensors="pt")
    summary_ids = model.generate(inputs["input_ids"][:, 1:], num_beams=4, min_length=20, max_length=100)
    summary = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
    return summary[0]


def greet(name):
    return "Hello " + name + "!!"


#TODO:
#  1. 下拉框，选择类目。   gr.Radio(['服饰','箱包', '鞋靴']
#  2. 支持NER、LM、Corrector
# beam seach参数
# promp参数

iface = gr.Interface(
    fn=generate,
    inputs=gr.inputs.Textbox(
        label="text",
        default="美的对开门风冷无霜家用智能电冰箱波光金纤薄机身高颜值助力保鲜，美的家居风，尺寸说明："
                "M以上的距离尤其是左右两侧距离必须保证。关于尺寸的更多问题可，LED冷光源，纤薄机身，风冷"
                "无霜，智能操控，远程调温，节能静音，照亮你的视野，535L大容量，系统散热和使用的便利性，"
                "建议左右两侧、顶部和背部需要预留10C，电源线和调平脚等。冰箱放置时为保证，菜谱推荐，半开"
                "门俯视图，全开门俯视图，预留参考图"),
    outputs=gr.Textbox(
            label="Summarization（文本摘要）",
        ),
    title="Abstractive Summarization (电商文本摘要)",
    description="电商领域文本摘要"
)
iface.launch()