Update app.py

app.py

@@ -1,12 +1,13 @@
 import gradio as gr
 from PyPDF2 import PdfReader
-from transformers import pipeline
+from transformers import pipeline, AutoModelForSeq2SeqLM, AutoTokenizer
 from gtts import gTTS
 from io import BytesIO
 import re
-import os
 
-summarizer = pipeline("summarization")
+# Load the LED-large model for summarization
+model_name = "pszemraj/led-large-book-summary"
+summarizer = pipeline("summarization", model=model_name, tokenizer=model_name)
 
 def extract_abstract_and_summarize(pdf_file):
     try:
@@ -27,17 +28,8 @@ def extract_abstract_and_summarize(pdf_file):
                     abstract_text = text[start_index:end_index]
                     break
 
-            # Summarize the extracted abstract
-            result = summarizer(
-                abstract_text,
-                min_length=16,
-                max_length=256,
-                no_repeat_ngram_size=3,
-                encoder_no_repeat_ngram_size=3,
-                repetition_penalty=3.5,
-                num_beams=4,
-                early_stopping=True,
-            )
+            # Summarize the extracted abstract using the LED-large model
+            result = summarizer(abstract_text, max_length=256, min_length=16, length_penalty=2.0)
             summary = result[0]['summary']
 
             # Generate audio
@@ -56,8 +48,7 @@ interface = gr.Interface(
     inputs=[gr.File(label="Upload PDF")],
     outputs=[gr.Textbox(label="Summary"), gr.Audio()],
     title="PDF Summarization & Audio Tool",
-    description="""PDF Summarization App. This app extracts the abstract from a PDF, summarizes it using the 'summarizer' model, and generates an audio of it. Only upload PDFs with abstracts. Please read the README.MD for information about the app and sample PDFs."""
+    description="""PDF Summarization App. This app extracts the abstract from a PDF, summarizes it using the 'pszemraj/led-large-book-summary' model, and generates an audio of it. Only upload PDFs with abstracts. Please read the README.MD for information about the app and sample PDFs."""
 )
 
 interface.launch()
-