Update model.py

model.py

@@ -18,9 +18,9 @@ from tqdm.notebook import tqdm
 class Abstractive_Summarization_Model:
     def __init__(self):
         self.text = None
-        self.IS_CNNDM = True # whether to use CNNDM dataset or XSum dataset
+        self.IS_CNNDM = False # whether to use CNNDM dataset or XSum dataset
         self.LOWER = False
-        self.max_length = 1024 if self.IS_CNNDM else 512
+        self.max_length = 1024 #if self.IS_CNNDM else 512
         self.model, self.tokenizer  = self.load_model()
         
     def load_model(self):
@@ -30,8 +30,8 @@ class Abstractive_Summarization_Model:
             model = BartForConditionalGeneration.from_pretrained('Yale-LILY/brio-cnndm-uncased')
             tokenizer = BartTokenizer.from_pretrained('Yale-LILY/brio-cnndm-uncased')
         else:
-            model = PegasusForConditionalGeneration.from_pretrained('Yale-LILY/brio-xsum-cased')
-            tokenizer = PegasusTokenizer.from_pretrained('Yale-LILY/brio-xsum-cased')
+            model = PegasusForConditionalGeneration.from_pretrained('tuner007/pegasus_paraphrase')
+            tokenizer = PegasusTokenizer.from_pretrained('tuner007/pegasus_paraphrase')
         print('[INFO]: Model Successfully Loaded :)')    
         return model, tokenizer