Update app.py

app.py

@@ -9,18 +9,16 @@ from pdfminer.high_level import extract_text
 from docx import Document
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
+from concurrent.futures import ThreadPoolExecutor
 
-# Load the models and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
 model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
 
 # Function to chunk text into sentence-based chunks
 def chunk_text(text, max_token_len=1024):
-    # Split text into sentences
     sentences = [sent.strip() + '.' for sent in re.split(r'(?<!\d)\.\s', text) if len(sent) > 1]
     token_lengths = [len(tokenizer.tokenize(sent)) for sent in sentences]
     
-    # Initialize chunking
     chunk_size = max_token_len
     chunks = []
     current_chunk = []
@@ -40,46 +38,30 @@ def chunk_text(text, max_token_len=1024):
     
     return chunks
 
-def docx_to_pdf(docx_file, output_pdf="converted_doc.pdf"):
-    doc = Document(docx_file)
-    full_text = []
-    for para in doc.paragraphs:
-        full_text.append(para.text)
-    
-    pdf = canvas.Canvas(output_pdf, pagesize=letter)
-    pdf.setFont("Helvetica", 12)
-    
-    text = pdf.beginText(40, 750)
-    for line in full_text:
-        text.textLine(line)
-    
-    pdf.drawText(text)
-    pdf.save()
-    return output_pdf
+# Summarization function
+def summarize_chunk(chunk, min_length=80):
+    inputs = tokenizer([chunk], max_length=1024, truncation=True, return_tensors="pt")
+    summary_ids = model.generate(inputs["input_ids"], num_beams=1, min_length=min_length, max_length=min_length + 300)
+    return tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
 
-# Main processing function using sentence-based chunking
+# Main processing function using parallel summarization
 def pdf_to_text(text, PDF, min_length=80):
     try:
         file_extension = os.path.splitext(PDF.name)[1].lower()
         
-        # If DOCX, convert to PDF
         if file_extension == '.docx':
             pdf_file_path = docx_to_pdf(PDF.name)
             text = extract_text(pdf_file_path)
-        # If PDF, extract text
         elif file_extension == '.pdf' and text == "":
             text = extract_text(PDF.name)
         
-        # Split text into chunks based on sentence boundaries
         chunks = chunk_text(text)
         summarized_text = ""
 
-        for chunk in chunks:
-            inputs = tokenizer([chunk], max_length=1024, truncation=True, return_tensors="pt")
-            min_length = int(min_length)
-            summary_ids = model.generate(inputs["input_ids"], num_beams=2, min_length=min_length, max_length=min_length + 400)
-            output_text = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
-            summarized_text += output_text + "\n\n"
+        # Parallelize summarization using ThreadPoolExecutor
+        with ThreadPoolExecutor() as executor:
+            summaries = list(executor.map(lambda chunk: summarize_chunk(chunk, min_length), chunks))
+            summarized_text = "\n\n".join(summaries)
 
         # Save summarized text to PDF
         pdf = FPDF()
@@ -99,7 +81,6 @@ def pdf_to_text(text, PDF, min_length=80):
     except Exception as e:
         return None, f"An error occurred: {str(e)}", None
 
-# Preloaded document processor
 def process_sample_document(min_length=80):
     sample_document_path = "Marbury v. Madison.pdf"
     
@@ -109,7 +90,7 @@ def process_sample_document(min_length=80):
 # Gradio interface
 with gr.Blocks() as iface:
     with gr.Row():
-        process_sample_button = gr.Button("Summarize Marbury v. Madison Case Pre-Uploaded")
+        process_sample_button = gr.Button("Summarize Pre-Uploaded Marbury v. Madison Case Document")
     
     text_input = gr.Textbox(label="Input Text")
     file_input = gr.File(label="Upload PDF or DOCX")