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

</br>

'[对话者0]<->汤姆 [对话者1]<->杰克'