Views
No views yet
1from transformers import MBartTokenizer, MBartForConditionalGeneration
2
3model_name = "IlyaGusev/mbart_ru_sum_gazeta"
4tokenizer = MBartTokenizer.from_pretrained(model_name)
5model = MBartForConditionalGeneration.from_pretrained(model_name)
6
7article_text = "..."
8
9input_ids = tokenizer(
10 [article_text],
11 max_length=600,
12 padding="max_length",
13 truncation=True,
14 return_tensors="pt",
15)["input_ids"]
16
17output_ids = model.generate(
18 input_ids=input_ids,
19 no_repeat_ngram_size=4
20)[0]
21
22summary = tokenizer.decode(output_ids, skip_special_tokens=True)
23print(summary)| Model | R-1-f | R-2-f | R-L-f | chrF | METEOR | BLEU | Avg char length |
|---|---|---|---|---|---|---|---|
| mbart_ru_sum_gazeta | 32.4 | 14.3 | 28.0 | 39.7 | 26.4 | 12.1 | 371 |
| rut5_base_sum_gazeta | 32.2 | 14.4 | 28.1 | 39.8 | 25.7 | 12.3 | 330 |
| rugpt3medium_sum_gazeta | 26.2 | 7.7 | 21.7 | 33.8 | 18.2 | 4.3 | 244 |
| Model | R-1-f | R-2-f | R-L-f | chrF | METEOR | BLEU | Avg char length |
|---|---|---|---|---|---|---|---|
| mbart_ru_sum_gazeta | 28.7 | 11.1 | 24.4 | 37.3 | 22.7 | 9.4 | 373 |
| rut5_base_sum_gazeta | 28.6 | 11.1 | 24.5 | 37.2 | 22.0 | 9.4 | 331 |
| rugpt3medium_sum_gazeta | 24.1 | 6.5 | 19.8 | 32.1 | 16.3 | 3.6 | 242 |
1import json
2import torch
3from transformers import MBartTokenizer, MBartForConditionalGeneration
4from datasets import load_dataset
5
6
7def gen_batch(inputs, batch_size):
8 batch_start = 0
9 while batch_start < len(inputs):
10 yield inputs[batch_start: batch_start + batch_size]
11 batch_start += batch_size
12
13
14def predict(
15 model_name,
16 input_records,
17 output_file,
18 max_source_tokens_count=600,
19 batch_size=4
20):
21 device = "cuda" if torch.cuda.is_available() else "cpu"
22
23 tokenizer = MBartTokenizer.from_pretrained(model_name)
24 model = MBartForConditionalGeneration.from_pretrained(model_name).to(device)
25
26 predictions = []
27 for batch in gen_batch(inputs, batch_size):
28 texts = [r["text"] for r in batch]
29 input_ids = tokenizer(
30 batch,
31 return_tensors="pt",
32 padding="max_length",
33 truncation=True,
34 max_length=max_source_tokens_count
35 )["input_ids"].to(device)
36
37 output_ids = model.generate(
38 input_ids=input_ids,
39 no_repeat_ngram_size=4
40 )
41 summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
42 for s in summaries:
43 print(s)
44 predictions.extend(summaries)
45 with open(output_file, "w") as w:
46 for p in predictions:
47 w.write(p.strip().replace("\n", " ") + "\n")
48
49gazeta_test = load_dataset('IlyaGusev/gazeta', script_version="v1.0")["test"]
50predict("IlyaGusev/mbart_ru_sum_gazeta", list(gazeta_test), "mbart_predictions.txt")1@InProceedings{10.1007/978-3-030-59082-6_9,
2 author="Gusev, Ilya",
3 editor="Filchenkov, Andrey and Kauttonen, Janne and Pivovarova, Lidia",
4 title="Dataset for Automatic Summarization of Russian News",
5 booktitle="Artificial Intelligence and Natural Language",
6 year="2020",
7 publisher="Springer International Publishing",
8 address="Cham",
9 pages="122--134",
10 isbn="978-3-030-59082-6"
11}