Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from docx import Document
 from io import BytesIO
+import re
 
 st.set_page_config(page_title="Doc Cleaner", layout="wide")
 
@@ -16,23 +17,26 @@ if uploaded_file:
 
     for para in doc.paragraphs:
         text = para.text.strip()
-        style_name = para.style.name.lower()
 
         if not text:
             continue
 
-        # 🔥 Improved heading detection
+        style_name = para.style.name.lower()
+
+        # ✅ Detection logic
         is_docx_heading = "heading" in style_name
         is_markdown_heading = text.startswith("#")
+        is_segment_heading = re.match(r"(segment\s*\d+)", text.lower())
+        is_colon_title = text.endswith(":") and len(text.split()) <= 6  # short titles like "Introduction:"
 
-        if is_docx_heading or is_markdown_heading:
+        if is_docx_heading or is_markdown_heading or is_segment_heading or is_colon_title:
             removed_headings.append(text)
         else:
             cleaned_paragraphs.append(text)
 
     cleaned_text = "\n\n".join(cleaned_paragraphs)
 
-    # ✅ Better UI
+    # 🗑️ Removed Headings Section
     st.subheader("🗑️ Removed Headings")
 
     if removed_headings:
@@ -40,6 +44,7 @@ if uploaded_file:
     else:
         st.info("No headings found")
 
+    # 📝 Cleaned Text Section
     st.subheader("📝 Cleaned Text")
 
     st.text_area(
@@ -48,13 +53,12 @@ if uploaded_file:
         height=400
     )
 
-    # ✅ Download (no button nesting issue)
-    buffer = BytesIO()
+    # 📥 Download cleaned DOCX
     new_doc = Document()
-
     for para in cleaned_paragraphs:
         new_doc.add_paragraph(para)
 
+    buffer = BytesIO()
     new_doc.save(buffer)
     buffer.seek(0)