svjack / summary-dialogue

from transformers import T5ForConditionalGeneration
from transformers import T5TokenizerFast as T5Tokenizer
model = "svjack/summary-dialogue"
device = "cpu"
tokenizer = T5Tokenizer.from_pretrained(model)
model = T5ForConditionalGeneration.from_pretrained(model).to(device).eval()
prompt =  "巴伐利亚号战列舰[a]（德语：SMS Bayern[b]）是德意志帝国海军巴伐利亚级战列舰的主导舰。该舰于1915年2月下水并于1916年7月开始服役，但已来不及参加日德兰海战。它的主炮包括分布在四座双联装炮塔中的八门380毫米口径炮，这比其前身国王级配备的十门305毫米口径炮有了显著改进。[c]舰只连同它的三艘姊妹舰已经形成了公海舰队第四战列分舰队的核心。而这当中仅有一艘，即巴登号完成建造；另外两艘则在第一次世界大战后期，当生产需求被转移至U型潜艇后而撤销。"
prompt = "摘要：{} 候选集：杰克 安娜".format(prompt)

encode = tokenizer(prompt, return_tensors='pt').to(device)
answer = model.generate(encode.input_ids,
                       max_length = 128,
    num_beams=2,
    top_p = 0.95,
    top_k = 50,
    repetition_penalty = 2.5,
    length_penalty=1.0,
    early_stopping=True,
                       )[0]
decoded = tokenizer.decode(answer, skip_special_tokens=True)
decoded.replace("安娜:", "\n").replace("杰克:", "\n").split("\n")

['',
 '巴罗利亚号战列舰是哪个国家? ',
 '德意志帝国海军的。它在1915年2月下水,1916年7月开始服役。 ',
 '该舰的主要装备是什么? ',
 '主炮包括四座双联装炮塔中的八门380毫米口径炮。 ',
 '这比其前身国王级装备的十门305毫米口径炮有了明显改进。 ',
 '但只有三艘姊妹舰已经形成公海舰队第四战列分舰队的核心。 ',
 '这是为什么? ',
 '它是二战后建造的。 ',
 '']