Update app.py

app.py

@@ -20,15 +20,15 @@ def summarize_text(text):
     for i, sentence in enumerate(sentences):
         if "Abstract" in sentence:
             start = i + 1
-            end = start + 4
+            end = start + 5
             break
             
     # Extract abstract 
     abstract = ". ".join(sentences[start:end+1])
 
     # Load BART model & tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("google/bigbird-pegasus-large-arxiv")
-    model = AutoModelForSeq2SeqLM.from_pretrained("google/bigbird-pegasus-large-arxiv")
+    tokenizer = AutoTokenizer.from_pretrained("pszemraj/led-base-book-summary")
+    model = AutoModelForSeq2SeqLM.from_pretrained(model="pszemraj/led-base-book-summary")
 
     # Tokenize abstract 
     inputs = tokenizer(abstract,
@@ -38,7 +38,7 @@ def summarize_text(text):
 
     # Generate summary
     summary_ids = model.generate(inputs['input_ids'],
-                                 num_beams=3,
+                                 #num_beams=3,
                                  max_length=50,
                                  min_length=30,
                                  do_sample=False,