Update app.py

app.py

@@ -1,5 +1,7 @@
 import gradio as gr
 import os
+import re
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from fpdf import FPDF
 from gtts import gTTS
@@ -8,20 +10,33 @@ from docx import Document
 from reportlab.lib.pagesizes import letter
 from reportlab.pdfgen import canvas
 
+# Load the models and tokenizer
 tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large-cnn")
 model = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
 
-# Function to split text into chunks based on token length
-def split_text_by_tokens(text, max_length=1024):
-    tokens = tokenizer.encode(text, return_tensors="pt", truncation=False)
-    total_length = tokens.shape[1]
+# 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 = []
+    current_length = 0
+    
+    for sent, length in zip(sentences, token_lengths):
+        if current_length + length <= chunk_size:
+            current_chunk.append(sent)
+            current_length += length
+        else:
+            chunks.append(" ".join(current_chunk))
+            current_chunk = [sent]
+            current_length = length
     
-    # Loop through the text, grabbing chunks of tokens
-    for i in range(0, total_length, max_length):
-        chunk_tokens = tokens[:, i:i+max_length]
-        chunk_text = tokenizer.decode(chunk_tokens[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
-        chunks.append(chunk_text)
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
     
     return chunks
 
@@ -42,18 +57,21 @@ def docx_to_pdf(docx_file, output_pdf="converted_doc.pdf"):
     pdf.save()
     return output_pdf
 
-# Main processing function with token-based text chunking
+# Main processing function using sentence-based chunking
 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)
         
-        chunks = split_text_by_tokens(text)
+        # Split text into chunks based on sentence boundaries
+        chunks = chunk_text(text)
         summarized_text = ""
 
         for chunk in chunks:
@@ -63,6 +81,7 @@ def pdf_to_text(text, PDF, min_length=80):
             output_text = tokenizer.batch_decode(summary_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[0]
             summarized_text += output_text + "\n\n"
 
+        # Save summarized text to PDF
         pdf = FPDF()
         pdf.add_page()
         pdf.set_font("Times", size=12)
@@ -70,6 +89,7 @@ def pdf_to_text(text, PDF, min_length=80):
         pdf_output_path = "legal.pdf"
         pdf.output(pdf_output_path)
         
+        # Convert summarized text to audio
         audio_output_path = "legal.wav"
         tts = gTTS(text=summarized_text, lang='en', slow=False)
         tts.save(audio_output_path)
@@ -79,12 +99,14 @@ 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"
     
     with open(sample_document_path, "rb") as f:
         return pdf_to_text("", f, min_length)
 
+# Gradio interface
 with gr.Blocks() as iface:
     with gr.Row():
         process_sample_button = gr.Button("Summarize Marbury v. Madison Case Pre-Uploaded")