improve workflow

app.py

@@ -20,141 +20,128 @@ def main():
         st.info("Please upload PDFs to continue.")
         return
 
-    # Extract text and keywords from PDFs
-    pdf_texts = {}
-    keywords_set = set()
-
-    # Initialize a progress bar
-    progress = st.progress(0)
-    total_files = len(uploaded_files)
-    processed_files = 0
-
-    st.info("Extracting keywords from PDFs...")
-    for uploaded_file in uploaded_files:
-        pdf_name = uploaded_file.name
-        try:
-            # Read PDF content
-            reader = PdfReader(uploaded_file)
-            text = ""
-            for page in reader.pages:
-                text += page.extract_text()
-            pdf_texts[pdf_name] = text.lower()
-
-            # Extract keywords using KeyBERT
-            extracted_keywords = kw_model.extract_keywords(text, top_n=5)
-            for kw, _ in extracted_keywords:
-                keywords_set.add(kw.lower())
-
-        except Exception as e:
-            st.error(f"Failed to process {pdf_name}: {e}")
-            continue
-
-        # Update progress
-        processed_files += 1
-        progress.progress(processed_files / total_files)
-
-    # Mark progress as complete
-    progress.progress(1.0)
+    # Check if uploaded files have changed
+    uploaded_file_names = [f.name for f in uploaded_files]
+    if "uploaded_files" not in st.session_state or st.session_state.uploaded_files != uploaded_file_names:
+        st.session_state.uploaded_files = uploaded_file_names
+        st.session_state.keywords_set = None
+
+    # Extract text and keywords from PDFs if not already done
+    if st.session_state.keywords_set is None:
+        st.info("Extracting keywords from PDFs...")
+        pdf_texts = {}
+        keywords_set = set()
+
+        progress1 = st.progress(0)
+        total_files = len(uploaded_files)
+
+        for i, uploaded_file in enumerate(uploaded_files):
+            pdf_name = uploaded_file.name
+            try:
+                reader = PdfReader(uploaded_file)
+                text = "".join(page.extract_text() for page in reader.pages)
+                pdf_texts[pdf_name] = text.lower()
+
+                extracted_keywords = kw_model.extract_keywords(text, top_n=5)
+                for kw, _ in extracted_keywords:
+                    keywords_set.add(kw.lower())
+            except Exception as e:
+                st.error(f"Failed to process {pdf_name}: {e}")
+            finally:
+                progress1.progress((i + 1) / total_files)
+
+        if not pdf_texts:
+            st.error("No PDFs could be processed.")
+            return
+
+        progress1.progress(1.0)
+        st.session_state.pdf_texts = pdf_texts
+        st.session_state.keywords_set = keywords_set
 
     # Display extracted keywords and let the user select topics
-    st.write("Extracted Keywords:")
     selected_keywords = st.multiselect(
         "Select at least two keywords/topics for grouping:",
-        list(keywords_set),
-        default=list(keywords_set)[:2]
+        list(st.session_state.keywords_set),
+        default=list(st.session_state.keywords_set)[:2]
     )
 
-    # Add a confirmation button to proceed
     if st.button("Confirm Keyword Selection"):
         if len(selected_keywords) < 2:
-            st.error("Please select at least two keywords.")
+            st.error("Please select at least two keywords to continue.")
         else:
-            st.success("Keyword selection confirmed. Processing PDFs...")
-            proceed_with_keywords = True
+            st.session_state.selected_keywords = selected_keywords
+            st.session_state.keywords_confirmed = True
     else:
-        proceed_with_keywords = False
+        st.session_state.keywords_confirmed = False
 
-    # Ensure the process does not continue until confirmation
-    if not proceed_with_keywords:
+    if not st.session_state.get("keywords_confirmed", False):
         st.stop()
 
-    # Initialize a progress bar for embedding precomputation
-    st.info("Precomputing embeddings for all PDFs...")
-    progress = st.progress(0)
-    total_pdfs = len(pdf_texts)
-    processed_pdfs = 0
+    st.success("Keyword selection confirmed. Processing PDFs...")
 
+    # Precompute embeddings for PDFs
+    st.info("Precomputing embeddings for PDFs...")
+    progress2 = st.progress(0)
     pdf_embeddings = {}
-    for pdf_name, text in pdf_texts.items():
+    pdf_texts = st.session_state.pdf_texts
+    total_pdfs = len(pdf_texts)
+
+    for i, (pdf_name, text) in enumerate(pdf_texts.items()):
         try:
-            # Compute embedding for the PDF
             pdf_embeddings[pdf_name] = semantic_model.encode(text, convert_to_tensor=True)
         except Exception as e:
             st.error(f"Failed to compute embedding for {pdf_name}: {e}")
-            continue
+        finally:
+            progress2.progress((i + 1) / total_pdfs)
 
-        # Update progress
-        processed_pdfs += 1
-        progress.progress(processed_pdfs / total_pdfs)
+    progress2.progress(1.0)
 
-    # Mark progress as complete
-    progress.progress(1.0)
-
-    # Initialize a progress bar for keyword embedding precomputation
+    # Precompute embeddings for selected keywords
     st.info("Precomputing embeddings for selected keywords...")
-    progress = st.progress(0)
+    progress3 = st.progress(0)
+    selected_keywords = st.session_state.selected_keywords
+    keyword_embeddings = {}
     total_keywords = len(selected_keywords)
-    processed_keywords = 0
 
-    keyword_embeddings = {}
-    for keyword in selected_keywords:
+    for i, keyword in enumerate(selected_keywords):
         try:
-            # Compute embedding for the keyword
             keyword_embeddings[keyword] = semantic_model.encode(keyword, convert_to_tensor=True)
         except Exception as e:
             st.error(f"Failed to compute embedding for keyword '{keyword}': {e}")
-            continue
-
-        # Update progress
-        processed_keywords += 1
-        progress.progress(processed_keywords / total_keywords)
+        finally:
+            progress3.progress((i + 1) / total_keywords)
 
-    # Mark progress as complete
-    progress.progress(1.0)
+    progress3.progress(1.0)
 
     # Group PDFs by the most relevant topic
+    st.info("Assigning PDFs to the most relevant topics...")
     pdf_groups = {keyword: [] for keyword in selected_keywords}
-    st.info("Assigning PDFs to the most relevant topic...")
 
     for pdf_name, text_embedding in pdf_embeddings.items():
-        max_similarity = -1
         best_keyword = None
+        max_similarity = -1
 
-        # Find the most similar keyword for this PDF
         for keyword, keyword_embedding in keyword_embeddings.items():
             similarity = util.pytorch_cos_sim(text_embedding, keyword_embedding).item()
             if similarity > max_similarity:
                 max_similarity = similarity
                 best_keyword = keyword
 
-        # Assign the PDF to the best matching keyword
         if best_keyword:
             pdf_groups[best_keyword].append(pdf_name)
 
     # Save grouped PDFs into folders
     output_folder = "grouped_pdfs"
     os.makedirs(output_folder, exist_ok=True)
+
     for keyword, pdf_names in pdf_groups.items():
         keyword_folder = os.path.join(output_folder, keyword)
         os.makedirs(keyword_folder, exist_ok=True)
         for pdf_name in pdf_names:
-            try:
-                matched_file = next(f for f in uploaded_files if f.name == pdf_name)
+            matched_file = next((f for f in uploaded_files if f.name == pdf_name), None)
+            if matched_file:
                 with open(os.path.join(keyword_folder, pdf_name), "wb") as f:
                     f.write(matched_file.getvalue())
-            except StopIteration:
-                st.error(f"File {pdf_name} not found in uploaded files.")
-                continue
 
     # Zip the folders
     zip_buffer = BytesIO()