Update README.md

README.md

@@ -104,6 +104,7 @@ Predicting all summaries:
 import json
 import torch
 from transformers import MBartTokenizer, MBartForConditionalGeneration
+from datasets import load_dataset
 
 
 def gen_batch(inputs, batch_size):
@@ -115,26 +116,19 @@ def gen_batch(inputs, batch_size):
 
 def predict(
     model_name,
-    test_file,
-    predictions_file,
-    targets_file,
+    input_records,
+    output_file,
     max_source_tokens_count=600,
-    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", " "))
-
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
     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):
+        texts = [r["text"] for r in batch]
         input_ids = tokenizer(
             batch,
             return_tensors="pt",
@@ -142,22 +136,21 @@ def predict(
             truncation=True,
             max_length=max_source_tokens_count
         )["input_ids"].to(device)
+        
         output_ids = model.generate(
             input_ids=input_ids,
-            repetition_penalty=3.0
+            no_repeat_ngram_size=4
         )
         summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         for s in summaries:
             print(s)
         predictions.extend(summaries)
-    with open(predictions_file, "w") as w:
+    with open(output_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")
+gazeta_test = load_dataset('IlyaGusev/gazeta', script_version="v1.0")["test"]
+predict("IlyaGusev/mbart_ru_sum_gazeta", list(gazeta_test), "mbart_predictions.txt")
 ```
 
 Evaluation: https://github.com/IlyaGusev/summarus/blob/master/evaluate.py