Update app.py

app.py

@@ -25,16 +25,19 @@ def summarize_text(text):
     for i, sentence in enumerate(sentences):
         if "Abstract" in sentence:
             start = i + 1
-            end = start + 3
+            end = start + 6
             break
     abstract = ". ".join(sentences[start:end+1]) 
 
   # Load BART model & tokenizer
-    tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
-    model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+    tokenizer = AutoTokenizer.from_pretrained("pszemraj/led-base-book-summary")
+    model = AutoModelForSeq2SeqLM.from_pretrained("pszemraj/led-base-book-summary")
 
     # Tokenize abstract 
-    inputs = tokenizer(abstract, return_tensors="pt", truncation=True)
+    inputs = tokenizer(abstract,
+                       max_length=512,
+                       return_tensors="pt",
+                       truncation=True)
 
     # Generate summary
     summary_ids = model.generate(inputs['input_ids'],