# 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
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 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]


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

"""
九阳电水壶热水壶家用电热水壶烧水壶304不锈钢快速加热容量开水壶无缝内胆九阳家用电水壶，贴心细节倾情陪你，沸腾仅需一首歌的时间，严选，以茶会友，早起，一杯柠檬水润肠道，食品级304不锈钢材质，赏心悦目厨房中的艺术品，无缝一体内胆，下班小聚，英国进口STRIX温控器，1.7L大容量，产品全家福，防烫保温杯体，材质，产品信息，可收纳底盘，加热开关，304不锈钢，大口径75°，广角开盖，给你一杯水的关怀，3-6秒确保，英国进口温控器，将多余电源线收纳在，更能确保每一杯水完全沸腾延时沸腾，不锈钢内胆，隔热层减缓热量散失速度，双层防烫，不锈钢滤网，自动断电，起到物理保温效果，不锈钢内盖，内部进行清洁，1800w超大加热底盘，能够让水快速沸腾，经久耐用，优质温控器12000次长使用寿命，碰撞简约优雅气息，满足一家人多种饮水需求，底盘，不锈钢壶
长虹2匹变频度精准控温立柜式冷暖空调白色节排，科技美，美，不止一点，健康美，8是，静音美，品质美，全自动，内外双静音安诸宁静，专利降噪，全铜管，0.1度精微感，层层过滤清新健康，美观大方，黑色加密滤网，美ei，宽电压启动，长虹空调所用钢材为攀钢与长虹共同研发多适用高密度钢板，由整张双面镀锌，一次冲压成型，没有焊接口，经过多次喷涂处理，使用寿命长。加密滤网，少风道内涡流现象，在不降低空调运行风量的前提下，还原室内外安湄。减少温度过调，轴流风扇，自动开机，带来的能源浪费，长虹空调的智能全自动，不仅能自动运作，同时空调上所具备的生态功能将，调节室温，省电恒温状态，选择模式，控制风速风向，低噪品牌压缩机，控温水平大大提升，新型风道技术，全部启动，创造理想家居环境。长虹空调
马兰士音响音箱家庭影院网络音频播放机双频段蓝牙黑色“乐”享新生活，Hi-Fi新时代，搭载高配置耳放，HEOs无线多房间音乐模块，精制优选元器件，广泛连接数字音乐，全面无线连接，远程控制，套装组合好搭档，氰听无损音乐，前后面板指示图，连接，输出，为原音而生与优雅相伴，全分离式电路，耳机放大器。精妙纯净，声场宽广。技术，尺寸，MarantzNA6006网络音频播放机，解码，声道，拓展，能连接计算机、外部硬盘、智能手机和平板电脑。镀金端口设计，电脑、外部硬盘、NAS、智能移动设备，高品质信号传输，前置USB-A端口可连接USB储存移动设备。音乐推送，模拟音频电路配备马兰士HDAM-SA2模块，低-中-高增益选择，匹配各种不同阻抗的耳机，灵活获得非凡体验。搭载HEOS无线多房间音乐模块可与所有支持该项技术的马兰士功放，宽440mm高106mm(
空气道壁挂式新风机系统家用除雾霾除甲醛白色净之本源，产品详情，噪音低至28分贝，键操控，智能辅热无惧严寒，让家中空气尽在掌握，EC电机2天仅耗1度电，0.5度，多种模式选择，五重纯物理过滤，大风量以一抵三，创新时代畅享新风30，空气道在民用新风领域率先采用H14级过滤材料对细颗粒物，二氧化碳浓度，去除有害微生物及过敏源>99%，PM2.5日均浓度，五重纯物理过滤系统将室外的新鲜空气经过物理过滤，后引入室内，无二次污染不产生对人体造成危害的臭，空气道起源于德国，为将先进的新风设计理念注入到空气道产品中，创，的单次通过滤除效率>99%，除去雾钼>99%，小时开机耗电仅0.5度*，绿色节能，氧，在增加室内含氧量的同时长效保护家人健康，在冬日也无惧室外严寒，在拥有室内清新空气，甲醛浓度
海尔8公斤节能静音高温消毒烫烫净全自动滚筒洗衣机靠实力说话，一掌控时间掌控自由，i-time智能时间洗，8公斤容量全家衣物一次清洗，细节绝不含糊，真正实力派，自动添加洗衣盒，洗羽绒服，就要专属程序，羊毛，牛仔，习绒，海尔洗衣机蓝晶系列滚筒，个性范儿，按照程序需求自动冲入洗衣机内，灵活旋钮，创新下排水洁净不残留，强力筋内筒，AMT防霉窗垫，LED大屏显示，洗衣液，消毒剂分别置放在洗衣盒中，从根本上解决污水残留问题避免，全新LD面板显示，更宽阔更大气操作信息一目了然，宽阔大气操作信息一目了然，右槽：消毒剂，简化洗衣程序，弹力筋中间的凹槽内分布，无残留排水模块，海尔洗衣机具有专业级羽绒洗护程序，为羽绒服营造洗护，一体化环境彻底告别手洗或者机洗，左槽：洗涤剂，我的智慧生活，中槽：柔顺剂，满足各种洗涤需求，告别昂贵洗衣店，自家
"""

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.inputs.Textbox(
        label="商品信息（product info）",
        default="美的对开门风冷无霜家用智能电冰箱波光金纤薄机身高颜值助力保鲜，美的家居风，尺寸说明："
                "M以上的距离尤其是左右两侧距离必须保证。关于尺寸的更多问题可，LED冷光源，纤薄机身，风冷"
                "无霜，智能操控，远程调温，节能静音，照亮你的视野，535L大容量，系统散热和使用的便利性，"
                "建议左右两侧、顶部和背部需要预留10C，电源线和调平脚等。冰箱放置时为保证，菜谱推荐，半开"
                "门俯视图，全开门俯视图，预留参考图"),
    outputs=gr.Textbox(
            label="文本摘要（Summarization）",
        ),
    examples=sum_examples,
    title="电商文本摘要（Abstractive Summarization）",
    description='电商领域文本摘要， 基于KPLUG预训练语言模型，'
                '<a href=""> K-PLUG: Knowledge-injected Pre-trained Language Model for Natural Language Understanding'
                ' and Generation in E-Commerce (Findings of EMNLP 2021) </a>。更多样例，见 https://github.com/xu-song/k-plug/tree/master/data_sample/sum/cepsum/jiadian/raw '
)

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