Update README.md

README.md

@@ -30,7 +30,7 @@ model_name = "IlyaGusev/rut5-base-sum-gazeta"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
-article_text = ".."
+article_text = "..."
 
 input_ids = tokenizer(
     [article_text],
@@ -68,7 +68,7 @@ Predicting all summaries:
 ```python
 import json
 import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import AutoTokenizer, T5ForConditionalGeneration
 from datasets import load_dataset
 
 
@@ -89,8 +89,8 @@ def predict(
 ):
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
-    tokenizer = MBartTokenizer.from_pretrained(model_name)
-    model = MBartForConditionalGeneration.from_pretrained(model_name).to(device)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name).to(device)
     
     predictions = []
     for batch in gen_batch(input_records, batch_size):
@@ -108,7 +108,6 @@ def predict(
             input_ids=input_ids,
             max_length=max_target_tokens_count,
             no_repeat_ngram_size=3,
-            num_beams=5,
             early_stopping=True
         )
         summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
@@ -120,7 +119,7 @@ def predict(
             w.write(p.strip().replace("\n", " ") + "\n")
 
 gazeta_test = load_dataset('IlyaGusev/gazeta', script_version="v1.0")["test"]
-predict("IlyaGusev/mbart_ru_sum_gazeta", gazeta_test["test"], "t5_predictions.txt")
+predict("IlyaGusev/rut5-base-sum-gazeta", list(gazeta_test), "t5_predictions.txt")
 ```
 
 Evaluation: https://github.com/IlyaGusev/summarus/blob/master/evaluate.py