README new sections

README.md

@@ -28,22 +28,22 @@ tokenizer = MBartTokenizer.from_pretrained(model_name)
 model = MBartForConditionalGeneration.from_pretrained(model_name)
 
 input_ids = tokenizer.prepare_seq2seq_batch(
-    [source],
-    src_lang="en_XX",
+    [article_text],
+    src_lang="en_XX", # fairseq training artifact
     return_tensors="pt",
     padding="max_length",
     truncation=True,
     max_length=600
-)["input_ids"][0]
+)["input_ids"]
 
 output_ids = model.generate(
-    input_ids=input_ids.unsqueeze(0),
+    input_ids=input_ids,
     max_length=162,
     no_repeat_ngram_size=3,
     num_beams=5,
-    top_k=0,
-    decoder_start_token_id=tokenizer.lang_code_to_id["ru_RU"]
+    top_k=0
 )[0]
+
 summary = tokenizer.decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
 print(summary)
 ```
@@ -64,21 +64,92 @@ print(summary)
 
 ## Eval results
 
+| Model                     | R-1-f | R-2-f | R-L-f | METEOR | BLEU |
+|:--------------------------|:------|:------|:------|:-------|:-----|
+| gazeta_mbart              | 32.6  | 14.6  | 28.2  | 25.7   | 49.8 |
+
+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"])
+
+    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() + "\n")
+    with open(targets_file, "w") as w:
+        for t in targets:
+            w.write(t.strip() + "\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"
+    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"
 }
 ```