|
--- |
|
language: |
|
- ru |
|
- ru-RU |
|
tags: |
|
- summarization |
|
- mbart |
|
datasets: |
|
- IlyaGusev/gazeta |
|
license: apache-2.0 |
|
--- |
|
|
|
# MBARTRuSumGazeta |
|
|
|
## Model description |
|
|
|
This is a ported version of [fairseq model](https://www.dropbox.com/s/fijtntnifbt9h0k/gazeta_mbart_v2_fairseq.tar.gz). |
|
|
|
For more details, please see, [Dataset for Automatic Summarization of Russian News](https://arxiv.org/abs/2006.11063). |
|
|
|
## Intended uses & limitations |
|
|
|
#### How to use |
|
|
|
```python |
|
from transformers import MBartTokenizer, MBartForConditionalGeneration |
|
|
|
article_text = "..." |
|
model_name = "IlyaGusev/mbart_ru_sum_gazeta" |
|
tokenizer = MBartTokenizer.from_pretrained(model_name) |
|
model = MBartForConditionalGeneration.from_pretrained(model_name) |
|
|
|
input_ids = tokenizer.prepare_seq2seq_batch( |
|
[article_text], |
|
src_lang="en_XX", # fairseq training artifact |
|
return_tensors="pt", |
|
padding="max_length", |
|
truncation=True, |
|
max_length=600 |
|
)["input_ids"] |
|
|
|
output_ids = model.generate( |
|
input_ids=input_ids, |
|
max_length=162, |
|
no_repeat_ngram_size=3, |
|
num_beams=5, |
|
top_k=0 |
|
)[0] |
|
|
|
summary = tokenizer.decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) |
|
print(summary) |
|
``` |
|
|
|
#### Limitations and bias |
|
|
|
- The model should work well with Gazeta.ru articles, but for any other agencies it can suffer from domain shift |
|
|
|
|
|
## Training data |
|
|
|
- Dataset: https://github.com/IlyaGusev/gazeta |
|
|
|
## Training procedure |
|
|
|
- Fairseq training script: https://github.com/IlyaGusev/summarus/blob/master/external/bart_scripts/train.sh |
|
- Porting: https://colab.research.google.com/drive/13jXOlCpArV-lm4jZQ0VgOpj6nFBYrLAr |
|
|
|
## Eval results |
|
|
|
| Model | R-1-f | R-2-f | R-L-f | METEOR | BLEU | |
|
|:--------------------------|:------|:------|:------|:-------|:-----| |
|
| gazeta_mbart | 32.3 | 14.3 | 27.9 | 25.5 | 12.4 | |
|
|
|
Predicting all summaries: |
|
```python |
|
import json |
|
import torch |
|
from transformers import MBartTokenizer, MBartForConditionalGeneration |
|
|
|
|
|
def gen_batch(inputs, batch_size): |
|
batch_start = 0 |
|
while batch_start < len(inputs): |
|
yield inputs[batch_start: batch_start + batch_size] |
|
batch_start += batch_size |
|
|
|
|
|
def predict( |
|
model_name, |
|
test_file, |
|
predictions_file, |
|
targets_file, |
|
max_source_tokens_count=600, |
|
max_target_tokens_count=160, |
|
use_cuda=True, |
|
batch_size=4 |
|
): |
|
inputs = [] |
|
targets = [] |
|
with open(test_file, "r") as r: |
|
for line in r: |
|
record = json.loads(line) |
|
inputs.append(record["text"]) |
|
targets.append(record["summary"].replace("\n", " ")) |
|
|
|
tokenizer = MBartTokenizer.from_pretrained(model_name) |
|
device = torch.device("cuda:0") if use_cuda else torch.device("cpu") |
|
model = MBartForConditionalGeneration.from_pretrained(model_name).to(device) |
|
predictions = [] |
|
for batch in gen_batch(inputs, batch_size): |
|
input_ids = tokenizer.prepare_seq2seq_batch( |
|
batch, |
|
src_lang="en_XX", |
|
return_tensors="pt", |
|
padding="max_length", |
|
truncation=True, |
|
max_length=max_source_tokens_count |
|
)["input_ids"].to(device) |
|
output_ids = model.generate( |
|
input_ids=input_ids, |
|
max_length=max_target_tokens_count + 2, |
|
no_repeat_ngram_size=3, |
|
num_beams=5, |
|
top_k=0 |
|
) |
|
summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False) |
|
for s in summaries: |
|
print(s) |
|
predictions.extend(summaries) |
|
with open(predictions_file, "w") as w: |
|
for p in predictions: |
|
w.write(p.strip().replace("\n", " ") + "\n") |
|
with open(targets_file, "w") as w: |
|
for t in targets: |
|
w.write(t.strip().replace("\n", " ") + "\n") |
|
|
|
predict("IlyaGusev/mbart_ru_sum_gazeta", "gazeta_test.jsonl", "predictions.txt", "targets.txt") |
|
``` |
|
|
|
Evaluation: https://github.com/IlyaGusev/summarus/blob/master/evaluate.py |
|
|
|
Flags: --language ru --tokenize-after --lower |
|
|
|
### BibTeX entry and citation info |
|
|
|
```bibtex |
|
@InProceedings{10.1007/978-3-030-59082-6_9, |
|
author="Gusev, Ilya", |
|
editor="Filchenkov, Andrey and Kauttonen, Janne and Pivovarova, Lidia", |
|
title="Dataset for Automatic Summarization of Russian News", |
|
booktitle="Artificial Intelligence and Natural Language", |
|
year="2020", |
|
publisher="Springer International Publishing", |
|
address="Cham", |
|
pages="122--134", |
|
isbn="978-3-030-59082-6" |
|
} |
|
``` |
|
|