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

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

import torch
import gradio as gr
from info import article
from kplug import modeling_kplug_s2s_patch
from transformers import BertTokenizer, BartForConditionalGeneration

model = BartForConditionalGeneration.from_pretrained("eson/kplug-base-cepsum-jiadian")  # cnn指的是cnn daily mail
tokenizer = BertTokenizer.from_pretrained("eson/kplug-base-cepsum-jiadian")

def summarize(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]

sum_examples = [
    "美的对开门风冷无霜家用智能电冰箱波光金纤薄机身高颜值助力保鲜，美的家居风，尺寸说明：M以上的距离尤其是左右两侧距离必须保证。关于尺寸的更多问题可，LED冷光源，纤薄机身，风冷无霜，智能操控，远程调温，节能静音，照亮你的视野，535L大容量，系统散热和使用的便利性，建议左右两侧、顶部和背部需要预留10C，电源线和调平脚等。冰箱放置时为保证，菜谱推荐，半开门俯视图，全开门俯视图，预留参考图",
    "爱家乐新加坡电风扇静音无叶风扇健康空气循环扇儿童球形风扇落地扇外观，宁静节能，产品结构，现代科技的结晶，品质，气家，未来风新时代，动里，空让，健康，低至13分贝/DC直流马达/低耗24，亲密玩伴，24W功率，/低耗，别加坡国民品牌，气流通道，增强室内空气运动，过尘栅网，1-12档风力调速，涡轮风扇，吸气口，大于6米随心掌控，电源适配暑，装箱明细，摆头角度，手动摇摆轨道，操作方式，与空调同时使用不仅可以让室温快速均衡作，电源插口，适用环境，还可以在短时间内，导引出风口，产品类型，快件重量，电机，暖空气向上冷空气向下，线长，使房间温度均衡，省电环保，定时，功率，将凉风或热风送给到附近的房间，轻松享受生活，左右自动(上下手动)摇摆9度，进风口，能够很快中和空气温度差",
    "海尔8公斤节能静音高温消毒烫烫净全自动滚筒洗衣机靠实力说话，一掌控时间掌控自由，i-time智能时间洗，8公斤容量全家衣物一次清洗，细节绝不含糊，真正实力派，自动添加洗衣盒，洗羽绒服，就要专属程序，羊毛，牛仔，习绒，海尔洗衣机蓝晶系列滚筒，个性范儿，按照程序需求自动冲入洗衣机内，灵活旋钮，创新下排水洁净不残留，强力筋内筒，AMT防霉窗垫，LED大屏显示，洗衣液，消毒剂分别置放在洗衣盒中，从根本上解决污水残留问题避免，全新LD面板显示，更宽阔更大气操作信息一目了然，宽阔大气操作信息一目了然，右槽：消毒剂，简化洗衣程序，弹力筋中间的凹槽内分布，无残留排水模块，海尔洗衣机具有专业级羽绒洗护程序，为羽绒服营造洗护，一体化环境彻底告别手洗或者机洗，左槽：洗涤剂，我的智慧生活，中槽：柔顺剂，满足各种洗涤需求，告别昂贵洗衣店，自家",
]

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

if __name__ == "__main__":
    sum_iface.launch()