mbart_ru_sum_gazeta / README.md
IlyaGusev's picture
Update README.md
37a48b2
|
raw
history blame
4.51 kB
---
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"
}
```