Update app.py

app.py

@@ -139,8 +139,9 @@ async def summarize_document(file, length="medium"):
     if pdf_path:
         result["pdfUrl"] = f"/files/{os.path.basename(pdf_path)}"
     return result"""
-from fastapi import FastAPI, UploadFile, File, Form
-from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
+# app.py
+
+from fastapi import UploadFile, File
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import fitz  # PyMuPDF
 import docx
@@ -158,23 +159,23 @@ import easyocr
 import datetime
 import hashlib
 
-# Initialize
+# Setup
 nltk.download('punkt', quiet=True)
-app = FastAPI()
 
-# Load Summarizer Model
+# 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)
 
-# Load OCR Reader
 reader = easyocr.Reader(['en'], gpu=torch.cuda.is_available())
 
-# Cache
 summary_cache = {}
 
+# Allowed file extensions
+ALLOWED_EXTENSIONS = {'pdf', 'docx', 'pptx', 'xlsx'}
+
 # --- Helper Functions ---
 
 def clean_text(text: str) -> str:
@@ -184,9 +185,9 @@ def clean_text(text: str) -> str:
     text = re.sub(r'\bPage\s*\d+\b', '', text, flags=re.IGNORECASE)
     return text.strip()
 
-def extract_text(file_path: str, file_extension: str):
+def extract_text(file_path: str, extension: str):
     try:
-        if file_extension == "pdf":
+        if extension == "pdf":
             with fitz.open(file_path) as doc:
                 text = "\n".join(page.get_text("text") for page in doc)
                 if len(text.strip()) < 50:
@@ -196,29 +197,24 @@ def extract_text(file_path: str, file_extension: str):
                     ocr_result = reader.readtext(temp_img.name, detail=0)
                     os.unlink(temp_img.name)
                     text = "\n".join(ocr_result) if ocr_result else text
-                return clean_text(text), ""
-
-        elif file_extension == "docx":
+        elif extension == "docx":
             doc = docx.Document(file_path)
-            return clean_text("\n".join(p.text for p in doc.paragraphs)), ""
-
-        elif file_extension == "pptx":
+            text = "\n".join(p.text for p in doc.paragraphs)
+        elif extension == "pptx":
             prs = pptx.Presentation(file_path)
-            text = [shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text")]
-            return clean_text("\n".join(text)), ""
-
-        elif file_extension == "xlsx":
+            text = "\n".join(shape.text for slide in prs.slides for shape in slide.shapes if hasattr(shape, "text"))
+        elif extension == "xlsx":
             wb = openpyxl.load_workbook(file_path, read_only=True)
-            text = [" ".join(str(cell) for cell in row if cell) for sheet in wb.sheetnames for row in wb[sheet].iter_rows(values_only=True)]
-            return clean_text("\n".join(text)), ""
-
-        elif file_extension in ["jpg", "jpeg", "png"]:
-            ocr_result = reader.readtext(file_path, detail=0)
-            return clean_text("\n".join(ocr_result)), ""
+            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 format."
+        
+        return clean_text(text), ""
 
-        return "", "Unsupported file format"
     except Exception as e:
-        return "", f"Error reading {file_extension.upper()} file: {str(e)}"
+        return "", f"Error reading {extension.upper()} file: {str(e)}"
 
 def chunk_text(text: str, max_tokens: int = 950):
     try:
@@ -243,7 +239,7 @@ def chunk_text(text: str, max_tokens: int = 950):
 
     return chunks
 
-def generate_summary(text: str, length: str = "medium") -> str:
+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]
@@ -283,14 +279,14 @@ def text_to_speech(text: str):
     except Exception:
         return ""
 
-def create_pdf(summary: str, original_filename: str):
+def create_pdf(summary: str, filename: str):
     try:
         pdf = FPDF()
         pdf.add_page()
         pdf.set_font("Arial", 'B', 16)
         pdf.cell(200, 10, txt="Document Summary", ln=1, align='C')
         pdf.set_font("Arial", size=12)
-        pdf.cell(200, 10, txt=f"Original file: {original_filename}", ln=1)
+        pdf.cell(200, 10, txt=f"Original file: {filename}", ln=1)
         pdf.cell(200, 10, txt=f"Generated on: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", ln=1)
         pdf.ln(10)
         pdf.multi_cell(0, 10, txt=summary)
@@ -300,49 +296,50 @@ def create_pdf(summary: str, original_filename: str):
     except Exception:
         return ""
 
-# --- API Endpoints ---
+# --- Public API Function ---
 
-@app.post("/summarize/")
-async def summarize_api(file: UploadFile = File(...), length: str = Form("medium")):
+async def summarize_document(file: UploadFile, length: str = "medium"):
     try:
-        contents = await file.read()
-        with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
-            tmp_file.write(contents)
+        filename = file.filename
+        extension = os.path.splitext(filename)[-1].lower().replace('.', '')
+
+        if extension not in ALLOWED_EXTENSIONS:
+            raise Exception(f"Unsupported file type: {extension.upper()}. Only PDF, DOCX, PPTX, XLSX are allowed.")
+
+        # Save uploaded file
+        with tempfile.NamedTemporaryFile(delete=False, suffix=f".{extension}") as tmp_file:
+            tmp_file.write(await file.read())
             tmp_path = tmp_file.name
 
-        file_ext = tmp_path.split('.')[-1].lower()
-        text, error = extract_text(tmp_path, file_ext)
+        # Extract text
+        text, error = extract_text(tmp_path, extension)
 
         if error:
-            return JSONResponse({"detail": error}, status_code=400)
+            os.unlink(tmp_path)
+            raise Exception(error)
 
         if not text or len(text.split()) < 30:
-            return JSONResponse({"detail": "Document too short to summarize"}, status_code=400)
+            os.unlink(tmp_path)
+            raise Exception("Document too short to summarize.")
 
+        # Summarize
         summary = generate_summary(text, length)
+
+        # Create audio + PDF
         audio_path = text_to_speech(summary)
-        pdf_path = create_pdf(summary, file.filename)
+        pdf_path = create_pdf(summary, filename)
 
+        # Clean temp file
+        os.unlink(tmp_path)
+
+        # Prepare response
         response = {"summary": summary}
         if audio_path:
             response["audioUrl"] = f"/files/{os.path.basename(audio_path)}"
         if pdf_path:
             response["pdfUrl"] = f"/files/{os.path.basename(pdf_path)}"
 
-        return JSONResponse(response)
+        return response
 
     except Exception as e:
-        print(f"Error during summarization: {str(e)}")
-        return JSONResponse({"detail": f"Internal server error: {str(e)}"}, status_code=500)
-
-@app.get("/files/{file_name}")
-async def serve_file(file_name: str):
-    path = os.path.join(tempfile.gettempdir(), file_name)
-    if os.path.exists(path):
-        return FileResponse(path)
-    return JSONResponse({"error": "File not found"}, status_code=404)
-
-@app.get("/")
-def home():
-    return RedirectResponse(url="/")
-
+        raise Exception(f"Summarization failed: {str(e)}")