Views
No views yet
1from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
2
3tokenizer = AutoTokenizer.from_pretrained("{repo_id}")
4model = AutoModelForSeq2SeqLM.from_pretrained("{repo_id}")
5
6dialogue = '''
7John: Hi Sarah, how are you doing?
8Sarah: I'm good, thanks for asking! Just finished my exams.
9John: That's great! How did they go?
10Sarah: I think I did well. I'm planning to celebrate tonight.
11'''
12
13inputs = tokenizer(dialogue, max_length=1024, return_tensors="pt", truncation=True)
14summary_ids = model.generate(inputs["input_ids"], num_beams=4, max_length=128, length_penalty=0.8)
15summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
16print(summary)