Update README.md

README.md

@@ -25,26 +25,25 @@ For more details, please see [Dataset for Automatic Summarization of Russian New
 ```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 = "..."
+input_ids = tokenizer(
     [article_text],
-    src_lang="en_XX", # fairseq training artifact
-    return_tensors="pt",
+    max_length=600,
     padding="max_length",
     truncation=True,
-    max_length=600
+    return_tensors="pt",
 )["input_ids"]
 
 output_ids = model.generate(
     input_ids=input_ids,
-    no_repeat_ngram_size=3
+    repetition_penalty=3.0
 )[0]
 
-summary = tokenizer.decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+summary = tokenizer.decode(output_ids, skip_special_tokens=True)
 print(summary)
 ```
 
@@ -55,12 +54,12 @@ print(summary)
 
 ## Training data
 
-- Dataset: https://github.com/IlyaGusev/gazeta
+- Dataset: [Gazeta](https://huggingface.co/datasets/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
+- Fairseq training script: [train.sh](https://github.com/IlyaGusev/summarus/blob/master/external/bart_scripts/train.sh)
+- Porting: [Colab link](https://colab.research.google.com/drive/13jXOlCpArV-lm4jZQ0VgOpj6nFBYrLAr)
 
 ## Eval results
 
@@ -98,7 +97,6 @@ def predict(
     predictions_file,
     targets_file,
     max_source_tokens_count=600,
-    max_target_tokens_count=160,
     use_cuda=True,
     batch_size=4
 ):
@@ -115,9 +113,8 @@ def predict(
     model = MBartForConditionalGeneration.from_pretrained(model_name).to(device)
     predictions = []
     for batch in gen_batch(inputs, batch_size):
-        input_ids = tokenizer.prepare_seq2seq_batch(
+        input_ids = tokenizer(
             batch,
-            src_lang="en_XX",
             return_tensors="pt",
             padding="max_length",
             truncation=True,
@@ -125,12 +122,9 @@ def predict(
         )["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
+            repetition_penalty=3.0
         )
-        summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        summaries = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         for s in summaries:
             print(s)
         predictions.extend(summaries)