Spaces:

darveen
/

text_summarizer

Runtime error

darveen commited on Oct 18, 2022

Commit

4115169

•

1 Parent(s): 4f8464b

Update model.py

Files changed (1) hide show

model.py CHANGED Viewed

@@ -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('tuner007/pegasus_paraphrase')
-            tokenizer = PegasusTokenizer.from_pretrained('tuner007/pegasus_paraphrase')
         print('[INFO]: Model Successfully Loaded :)')
         return model, tokenizer

             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')
         print('[INFO]: Model Successfully Loaded :)')
         return model, tokenizer