Update modules/abstractive.py

modules/abstractive.py

@@ -1,10 +1,9 @@
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration
 
 def load_summarizers():
     models = {
-        "Pegasus": "google/pegasus-large",
         "T5": "Overglitch/t5-small-cnn-dailymail",
         "BART": "facebook/bart-large-cnn",
     }
@@ -15,6 +14,11 @@ def load_summarizers():
         summarizers[model_name] = (model, tokenizer)
     return summarizers
 
+def load_pegasus_model_and_tokenizer(model_name: str):
+    model = PegasusForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = PegasusTokenizer.from_pretrained(model_name)
+    return model, tokenizer
+
 
 def abstractive_summary(summarizers, model_name, text, max_length, num_beams):
     model, tokenizer = summarizers[model_name]