Views
No views yet
1import torch
2from transformers import LEDTokenizer, LEDForConditionalGeneration
3tokenizer = LEDTokenizer.from_pretrained("bakhitovd/led-base-7168-ml")
4model = LEDForConditionalGeneration.from_pretrained("bakhitovd/led-base-7168-ml")1article = "... long document ..."
2inputs_dict = tokenizer.encode(article, padding="max_length", max_length=16384, return_tensors="pt", truncation=True)
3input_ids = inputs_dict.input_ids.to("cuda")
4attention_mask = inputs_dict.attention_mask.to("cuda")
5global_attention_mask = torch.zeros_like(attention_mask)
6global_attention_mask[:, 0] = 1
7predicted_abstract_ids = model.generate(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask, max_length=512)
8summary = tokenizer.decode(predicted_abstract_ids, skip_special_tokens=True)
9print(summary)