Update modules/abstractive.py

modules/abstractive.py

@@ -1,24 +1,25 @@
 import torch
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers import PegasusTokenizer, PegasusForConditionalGeneration
+from transformers import PegasusTokenizer, PegasusForConditionalGeneration 
 
 def load_summarizers():
     models = {
+        "Pegasus": "google/pegasus-cnn_dailymail"
         "T5": "Overglitch/t5-small-cnn-dailymail",
         "BART": "facebook/bart-large-cnn",
     }
     summarizers = {}
     for model_name, model_path in models.items():
+        if model_name == "Pegasus":
+            tokenizer = PegasusTokenizer.from_pretrained(model_path)
+        else:
+            tokenizer = AutoTokenizer.from_pretrained(model_path)
+            
         model = AutoModelForSeq2SeqLM.from_pretrained(model_path).to("cuda" if torch.cuda.is_available() else "cpu")
         tokenizer = AutoTokenizer.from_pretrained(model_path)
         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]