Views
No views yet
1
2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3device = "cuda" if torch.cuda.is_available() else "CPU"
4
5checkpoint = "d0r1h/led-base-ilc"
6
7tokenizer = AutoTokenizer.from_pretrained(checkpoint)
8model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, return_dict_in_generate=True).to(device)
9case = "......."
10input_ids = tokenizer(case, return_tensors="pt").input_ids.to(device)
11global_attention_mask = torch.zeros_like(input_ids)
12global_attention_mask[:, 0] = 1
13sequences = model.generate(input_ids,
14 global_attention_mask=global_attention_mask).sequences
15summary = tokenizer.batch_decode(sequences,
16 skip_special_tokens=True)
17 | Model | rouge1-f | rouge1-p | rouge2-f | rouge2-p | rougeL-f | rougeL-p |
|---|---|---|---|---|---|---|
| led-ilc | 42 | 47 | 22 | 24 | 39 | 44 |
| led-base | 3 | 39 | 1 | 21 | 3 | 37 |