import sys
sys.path.append("./fengshen/examples/pegasus/")

import fengshen
import gradio as gr
from transformers import PegasusForConditionalGeneration
from tokenizers_pegasus import PegasusTokenizer

model = PegasusForConditionalGeneration.from_pretrained("IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese")
tokenizer = PegasusTokenizer.from_pretrained("IDEA-CCNL/Randeng-Pegasus-523M-Summary-Chinese")


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

#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
def generator(text, max_length, num_return_sequences):
  inputs = tokenizer(text, max_length=1024, return_tensors="pt")
  # Generate Summary
  summary_ids = model.generate(inputs["input_ids"], max_length)
  return tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[:num_return_sequences]

def generate(text):
    result = generator(text, max_length=300, num_return_sequences=1)
    return result

examples = [
    ["""现在乌军的迫击炮都可以开始轰击库皮扬斯克的俄军目标了，双方相距只有几公里。
最重要的是该地区俄军背后就是奥斯科尔河-北顿涅茨河。而为了确保第聂伯河上赫尔松城的后勤保障，俄军已经把舟桥部队主力调到赫尔松去了，现在是远水解不了近渴。
乌军很有可能暂时不会直接攻击伊久姆城区，现在还是要先拿下库皮扬斯克和奥斯科尔河上的两座桥梁。
切断俄军从别尔哥罗德向库皮扬斯克的物资输送，切断库皮扬斯克-伊久姆公路，断其粮道。
现在要看俄军有没有预备队。库皮扬斯克没援军的话大概率守不住。但是现在打了三天，俄军还没有援军抵达战场。
现在，俄军最主要任务是守住库皮扬斯克，同时要确保库皮扬斯克-伊久姆高速公路的安全。"""],
    ["在北京冬奥会自由式滑雪女子坡面障碍技巧决赛中，中国选手谷爱凌夺得银牌。祝贺谷爱凌！今天上午，自由式滑雪女子坡面障碍技巧决赛举行。决赛分三轮进行，取选手最佳成绩排名决出奖牌。第一跳，中国选手谷爱凌获得69.90分。在12位选手中排名第三。完成动作后，谷爱凌又扮了个鬼脸，甚是可爱。第二轮中，谷爱凌在道具区第三个障碍处失误，落地时摔倒。获得16.98分。网友：摔倒了也没关系，继续加油！在第二跳失误摔倒的情况下，谷爱凌顶住压力，第三跳稳稳发挥，流畅落地！获得86.23分！此轮比赛，共12位选手参赛，谷爱凌第10位出场。网友：看比赛时我比谷爱凌紧张，加油！"],
]

iface = gr.Interface(
    fn=generate,
    inputs=gr.inputs.Textbox(lines=5, label="Input Text"),
    outputs=gr.outputs.Textbox(label="Generated Text"),
    examples=examples
)
# 

iface.launch()