Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import gradio as gr
 import google.generativeai as genai
 from markdown_pdf import MarkdownPdf, Section
-import pikepdf
+import subprocess
 
 # ---------- PROMPTS ----------
 PROMPTS = {
@@ -104,24 +104,36 @@ def save_as_pdf(text, filename="output.pdf"):
     pdf.save(filename)
     return filename
 
-# ---------- HELPER: PDF Compression ----------
-def compress_pdf(input_path, output_path):
+# ---------- HELPER: Compress PDF ----------
+def compress_pdf(input_path, output_path=None, max_size=20*1024*1024):
+    """Compress PDF using Ghostscript if larger than max_size (default 20MB)."""
+    if output_path is None:
+        base, ext = os.path.splitext(input_path)
+        output_path = f"{base}_compressed{ext}"
+
+    if os.path.getsize(input_path) <= max_size:
+        return input_path  # No compression needed
+
     try:
-        pdf = pikepdf.open(input_path)
-        pdf.save(output_path, optimize_version=True)
-        pdf.close()
-        return output_path
+        gs_cmd = [
+            "gs", "-sDEVICE=pdfwrite",
+            "-dCompatibilityLevel=1.4",
+            "-dPDFSETTINGS=/ebook",   # options: /screen, /ebook, /printer
+            "-dNOPAUSE", "-dQUIET", "-dBATCH",
+            f"-sOutputFile={output_path}", input_path
+        ]
+        subprocess.run(gs_cmd, check=True)
+
+        if os.path.getsize(output_path) <= max_size:
+            print(f"✅ Compressed {input_path} → {output_path}")
+            return output_path
+        else:
+            print(f"⚠️ Compression failed to reduce below {max_size/1024/1024} MB")
+            return input_path
     except Exception as e:
-        print(f"❌ Compression failed for {input_path}: {e}")
+        print(f"⚠️ Compression error: {e}")
         return input_path
 
-def check_and_compress(path):
-    if os.path.getsize(path) > 20 * 1024 * 1024:  # 20 MB
-        compressed_path = os.path.splitext(path)[0] + "_compressed.pdf"
-        print(f"⚡ Compressing {os.path.basename(path)} (>20MB)...")
-        return compress_pdf(path, compressed_path)
-    return path
-
 # ---------- HELPER: Create Model with Fallback ----------
 def create_model():
     try:
@@ -134,19 +146,19 @@ def create_model():
 # ---------- PIPELINE: ALIGN + GRADE ----------
 def align_and_grade(qp_file, ms_file, ans_file):
     try:
-        # Ensure files are <20MB
-        qp_file = check_and_compress(qp_file)
-        ms_file = check_and_compress(ms_file)
-        ans_file = check_and_compress(ans_file)
+        # Step 0: Compress if needed
+        qp_file = compress_pdf(qp_file, "qp_compressed.pdf")
+        ms_file = compress_pdf(ms_file, "ms_compressed.pdf")
+        ans_file = compress_pdf(ans_file, "ans_compressed.pdf")
 
-        # Uploads
+        # Step 1: Uploads
         qp_uploaded = genai.upload_file(path=qp_file, display_name="Question Paper")
         ms_uploaded = genai.upload_file(path=ms_file, display_name="Markscheme")
         ans_uploaded = genai.upload_file(path=ans_file, display_name="Answer Sheet")
 
         model = create_model()
 
-        # Step 1: Alignment
+        # Step 2: Alignment
         resp = model.generate_content([
             PROMPTS["ALIGNMENT_PROMPT"]["content"],
             qp_uploaded,
@@ -159,7 +171,7 @@ def align_and_grade(qp_file, ms_file, ans_file):
 
         aligned_pdf_path = save_as_pdf(aligned_text, "aligned_qp_ms_as.pdf")
 
-        # Step 2: Grading (automatic)
+        # Step 3: Grading
         response = model.generate_content([
             PROMPTS["GRADING_PROMPT"]["content"],
             aligned_text