from transformers import T5ForConditionalGeneration
from transformers import T5TokenizerFast as T5Tokenizer
model = "svjack/dialogue-summary"
device = "cpu"
tokenizer = T5Tokenizer.from_pretrained(model)
model = T5ForConditionalGeneration.from_pretrained(model).to(device).eval()
prompt =  '''
汤姆：能告诉我如何去机场吗？
杰克：我认为你最好座机场大巴。
汤姆：可是我有些晕车。
杰克：如果你觉得颠簸的话，也可以坐地铁去。
汤姆：地铁需要多长时间？
杰克：1小时左右。
汤姆：周围有什么好的风景吗？
杰克：路过龙井山。
汤姆：那里有哪些特产？
杰克：山中出产龙井茶。
'''

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

''' '汤姆想坐地铁去机场。杰克建议他到龙井山,那里有龙井茶。' '''