Views
No views yet
1from transformers import T5ForConditionalGeneration
2from transformers import T5TokenizerFast as T5Tokenizer
3model = "svjack/dialogue-summary-fill-characters"
4device = "cpu"
5tokenizer = T5Tokenizer.from_pretrained(model)
6model = T5ForConditionalGeneration.from_pretrained(model).to(device).eval()
7
8x = '''
9 根据上下文填补人物映射：
10 上下文：汤姆：能告诉我如何去机场吗？
11杰克：我认为你最好座机场大巴。
12汤姆：可是我有些晕车。
13杰克：如果你觉得颠簸的话，也可以坐地铁去。
14汤姆：地铁需要多长时间？
15杰克：1小时左右。
16汤姆：周围有什么好的风景吗？
17杰克：路过龙井山。
18汤姆：那里有哪些特产？
19杰克：山中出产龙井茶。
20 摘要：[对话者0]问去机场的交通工具，[对话者1]推荐地铁，路过龙井山
21 候选集： 杰克 汤姆
22 答案：
23'''
24
25encode = tokenizer(x, return_tensors='pt').to(device)
26answer = model.generate(encode.input_ids,
27 max_length = 128,
28 num_beams=2,
29 top_p = 0.95,
30 top_k = 50,
31 repetition_penalty = 2.5,
32 length_penalty=1.0,
33 early_stopping=True,
34 )[0]
35decoded = tokenizer.decode(answer, skip_special_tokens=True)
36decoded'[对话者0]<->汤姆 [对话者1]<->杰克'