Update app.py

app.py

@@ -1,57 +1,57 @@
-# app_logic.py
+# app.py
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
-import fitz, docx, pptx, openpyxl, re, nltk, tempfile, os, easyocr, hashlib, datetime
+import fitz, docx, pptx, openpyxl, re, nltk, tempfile, os, easyocr, datetime, hashlib
 from nltk.tokenize import sent_tokenize
 from fpdf import FPDF
 from gtts import gTTS
 
 nltk.download('punkt', quiet=True)
 
-# Load once
+# Load models
 MODEL_NAME = "facebook/bart-large-cnn"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
+model.eval()
 summarizer = pipeline("summarization", model=model, tokenizer=tokenizer, device=-1, batch_size=4)
 reader = easyocr.Reader(['en'], gpu=False)
 
 summary_cache = {}
 
-def clean_text(text):
+def clean_text(text: str) -> str:
     text = re.sub(r'\s+', ' ', text)
     text = re.sub(r'\u2022\s*|\d\.\s+', '', text)
     text = re.sub(r'\[.*?\]|\(.*?\)', '', text)
     text = re.sub(r'\bPage\s*\d+\b', '', text, flags=re.IGNORECASE)
     return text.strip()
 
-def extract_text(file_path, file_extension):
+def extract_text(file_path: str, ext: str):
     try:
-        if file_extension in ["pdf"]:
+        if ext == "pdf":
             with fitz.open(file_path) as doc:
                 text = "\n".join(page.get_text("text") for page in doc)
                 if len(text.strip()) < 50:
                     images = [page.get_pixmap() for page in doc]
                     temp_img = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
                     images[0].save(temp_img.name)
-                    ocr_result = reader.readtext(temp_img.name, detail=0)
+                    text = "\n".join(reader.readtext(temp_img.name, detail=0))
                     os.unlink(temp_img.name)
-                    text = "\n".join(ocr_result) if ocr_result else text
-        elif file_extension in ["docx"]:
+        elif ext == "docx":
             doc = docx.Document(file_path)
             text = "\n".join(p.text for p in doc.paragraphs)
-        elif file_extension in ["pptx"]:
+        elif ext == "pptx":
             prs = pptx.Presentation(file_path)
             text = "\n".join(shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text"))
-        elif file_extension in ["xlsx"]:
+        elif ext == "xlsx":
             wb = openpyxl.load_workbook(file_path, read_only=True)
             text = "\n".join([" ".join(str(cell) for cell in row if cell) for sheet in wb.sheetnames for row in wb[sheet].iter_rows(values_only=True)])
         else:
-            return "", "Unsupported file type"
-        
-        return clean_text(text), ""
+            text = ""
     except Exception as e:
-        return "", f"Extraction error: {e}"
+        return "", f"Error extracting text: {str(e)}"
 
-def chunk_text(text, max_tokens=950):
+    return clean_text(text), ""
+
+def chunk_text(text: str, max_tokens: int = 950):
     sentences = sent_tokenize(text)
     chunks, current_chunk = [], ""
     for sentence in sentences:
@@ -64,21 +64,36 @@ def chunk_text(text, max_tokens=950):
         chunks.append(current_chunk.strip())
     return chunks
 
-def generate_summary(text, length="medium"):
+def generate_summary(text: str, length: str = "medium"):
     cache_key = hashlib.md5((text + length).encode()).hexdigest()
     if cache_key in summary_cache:
         return summary_cache[cache_key]
 
-    params = {"short": (30, 80), "medium": (80, 200), "long": (210, 300)}[length]
-    min_len, max_len = params
+    length_params = {
+        "short": {"max_length": 80, "min_length": 30},
+        "medium": {"max_length": 200, "min_length": 80},
+        "long": {"max_length": 300, "min_length": 210}
+    }
 
     chunks = chunk_text(text)
-    summaries = summarizer(chunks, max_length=max_len, min_length=min_len, do_sample=False)
+    summaries = summarizer(
+        chunks,
+        max_length=length_params[length]["max_length"],
+        min_length=length_params[length]["min_length"],
+        do_sample=False,
+        truncation=True,
+        no_repeat_ngram_size=2,
+        num_beams=2,
+        early_stopping=True
+    )
     final_summary = " ".join(s['summary_text'] for s in summaries)
+    final_summary = ". ".join(s.strip().capitalize() for s in final_summary.split(". ") if s.strip())
+    final_summary = final_summary if len(final_summary) > 25 else "Summary too short."
+
     summary_cache[cache_key] = final_summary
     return final_summary
 
-def text_to_speech(text):
+def text_to_speech(text: str):
     try:
         tts = gTTS(text)
         temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
@@ -87,12 +102,10 @@ def text_to_speech(text):
     except:
         return ""
 
-def create_pdf(summary, original_filename):
+def create_pdf(summary: str, filename: str):
     try:
         pdf = FPDF()
         pdf.add_page()
-        pdf.set_font("Arial", 'B', 16)
-        pdf.cell(200, 10, "Summary", ln=True, align='C')
         pdf.set_font("Arial", size=12)
         pdf.multi_cell(0, 10, summary)
         temp_pdf = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
@@ -100,3 +113,29 @@ def create_pdf(summary, original_filename):
         return temp_pdf.name
     except:
         return ""
+
+async def summarize_document(file, length="medium"):
+    contents = await file.read()
+    with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
+        tmp_file.write(contents)
+        tmp_path = tmp_file.name
+
+    ext = file.filename.split('.')[-1].lower()
+    text, error = extract_text(tmp_path, ext)
+
+    if error:
+        raise Exception(error)
+
+    if not text or len(text.split()) < 30:
+        raise Exception("Document too short to summarize.")
+
+    summary = generate_summary(text, length)
+    audio_path = text_to_speech(summary)
+    pdf_path = create_pdf(summary, file.filename)
+
+    result = {"summary": summary}
+    if audio_path:
+        result["audioUrl"] = f"/files/{os.path.basename(audio_path)}"
+    if pdf_path:
+        result["pdfUrl"] = f"/files/{os.path.basename(pdf_path)}"
+    return result