update: main.py

main.py

@@ -3,9 +3,9 @@ from datetime import datetime
 import streamlit as st
 from io import BytesIO
 from fpdf import FPDF
-import time
-import requests
 import tempfile
+import requests
+import time
 
 # Configure the page
 st.set_page_config(
@@ -49,6 +49,27 @@ def get_dejavu_font_path():
     tempFontFile.close()
     return tempFontFile.name
 
+# Cache the summary generation
+@st.cache_data(show_spinner=False)
+def generate_summary(pipeline, pdf_bytes):
+    return pipeline.run(pdf_bytes)
+
+# Cache the PDF generation
+@st.cache_data(show_spinner=False)
+def generate_pdf_bytes(summary, font_path):
+    pdf = FPDF()
+    pdf.add_page()
+    pdf.add_font("DejaVu", "", font_path, uni=True)
+    pdf.set_font("DejaVu", '', 16)
+    pdf.cell(200, 10, txt="PDF Summary", ln=True, align='C')
+    pdf.ln(10)
+    pdf.set_font("DejaVu", '', 12)
+    pdf.multi_cell(0, 10, summary)
+    
+    pdf_output = BytesIO()
+    pdf.output(pdf_output)
+    return pdf_output.getvalue()
+
 # Sidebar
 with st.sidebar:
     st.markdown("### Upload PDF")
@@ -91,7 +112,7 @@ if uploadedFile:
             # Generate summary
             statusBox.info("Generating summary...")
             summaryStartTime = time.time()
-            summary = pipeline.run(pdfBytes)
+            summary = generate_summary(pipeline, pdfBytes)
             totalTime = time.time() - startTime
 
             if summary:
@@ -104,23 +125,12 @@ if uploadedFile:
                     if st.button("Download Summary as PDF"):
                         try:
                             font_path = get_dejavu_font_path()
-
-                            pdf = FPDF()
-                            pdf.add_page()
-                            pdf.add_font("DejaVu", "", font_path, uni=True)
-                            pdf.set_font("DejaVu", '', 16)
-                            pdf.cell(200, 10, txt="PDF Summary", ln=True, align='C')
-                            pdf.ln(10)
-                            pdf.set_font("DejaVu", '', 12)
-                            pdf.multi_cell(0, 10, summary)
-
-                            pdfOutput = BytesIO()
-                            pdf.output(pdfOutput)
+                            pdf_bytes = generate_pdf_bytes(summary, font_path)
                             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
 
                             st.download_button(
                                 label="Click to Download",
-                                data=pdfOutput.getvalue(),
+                                data=pdf_bytes,
                                 file_name=f"summary_{timestamp}.pdf",
                                 mime="application/pdf"
                             )