from transformers import T5ForConditionalGeneration
from transformers import T5TokenizerFast as T5Tokenizer
model = "svjack/dialogue-summary-fill-characters"
device = "cpu"
tokenizer = T5Tokenizer.from_pretrained(model)
model = T5ForConditionalGeneration.from_pretrained(model).to(device).eval()
x = '''
根据上下文填补人物映射：
上下文：汤姆：能告诉我如何去机场吗？
杰克：我认为你最好座机场大巴。
汤姆：可是我有些晕车。
杰克：如果你觉得颠簸的话，也可以坐地铁去。
汤姆：地铁需要多长时间？
杰克：1小时左右。
汤姆：周围有什么好的风景吗？
杰克：路过龙井山。
汤姆：那里有哪些特产？
杰克：山中出产龙井茶。
摘要：[对话者0]问去机场的交通工具，[对话者1]推荐地铁，路过龙井山
候选集： 杰克 汤姆
答案：
'''
encode = tokenizer(x, 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
'[对话者0]<->汤姆 [对话者1]<->杰克'
