Update app.py

app.py

@@ -18,11 +18,9 @@ logging.basicConfig(level=logging.INFO)
 MODEL_ID = "zai-org/GLM-OCR"
 OCR_API_PORT = 8080
 OCR_API_BASE = f"http://127.0.0.1:{OCR_API_PORT}"
-# Fixed runtime settings (no environment variables required).
 VLLM_MAX_MODEL_LEN = "8192"
 VLLM_GPU_MEMORY_UTIL = "0.75"
 VLLM_EXTRA_ARGS = ""
-# Keep output budget below context length to leave room for prompt tokens.
 GLMOCR_MAX_OUTPUT_TOKENS = 2048
 
 _parser = None
@@ -39,7 +37,8 @@ def _render_pdf_pages_to_images(pdf_path: str) -> List[str]:
             page = doc[i]
             pix = page.get_pixmap(matrix=fitz.Matrix(2.5, 2.5), alpha=False)
             img_path = os.path.join(
-                tempfile.gettempdir(), f"glmocr_page_{os.getpid()}_{i}_{int(time.time()*1000)}.png"
+                tempfile.gettempdir(),
+                f"glmocr_page_{os.getpid()}_{i}_{int(time.time()*1000)}.png",
             )
             pix.save(img_path)
             page_images.append(img_path)
@@ -54,19 +53,12 @@ def _start_vllm_if_needed() -> None:
         return
 
     cmd: List[str] = [
-        "python",
-        "-m",
-        "vllm.entrypoints.openai.api_server",
-        "--model",
-        MODEL_ID,
-        "--served-model-name",
-        "glm-ocr",
-        "--port",
-        str(OCR_API_PORT),
-        "--max-model-len",
-        VLLM_MAX_MODEL_LEN,
-        "--gpu-memory-utilization",
-        VLLM_GPU_MEMORY_UTIL,
+        "python", "-m", "vllm.entrypoints.openai.api_server",
+        "--model", MODEL_ID,
+        "--served-model-name", "glm-ocr",
+        "--port", str(OCR_API_PORT),
+        "--max-model-len", VLLM_MAX_MODEL_LEN,
+        "--gpu-memory-utilization", VLLM_GPU_MEMORY_UTIL,
     ]
     if VLLM_EXTRA_ARGS:
         cmd.extend(VLLM_EXTRA_ARGS.split())
@@ -98,10 +90,8 @@ def _configure_glmocr_to_keep_header_footer() -> None:
         cfg = yaml.safe_load(f) or {}
 
     pipeline = cfg.setdefault("pipeline", {})
-
     maas = pipeline.setdefault("maas", {})
     maas["enabled"] = False
-
     ocr_api = pipeline.setdefault("ocr_api", {})
     ocr_api["api_host"] = "127.0.0.1"
     ocr_api["api_port"] = OCR_API_PORT
@@ -111,23 +101,17 @@ def _configure_glmocr_to_keep_header_footer() -> None:
     ocr_api["verify_ssl"] = False
     page_loader = pipeline.setdefault("page_loader", {})
     page_loader["max_tokens"] = GLMOCR_MAX_OUTPUT_TOKENS
-
     layout = pipeline.setdefault("layout", {})
     label_task_mapping = layout.setdefault("label_task_mapping", {})
     text_labels = set(label_task_mapping.get("text", []) or [])
     abandon_labels = set(label_task_mapping.get("abandon", []) or [])
     skip_labels = set(label_task_mapping.get("skip", []) or [])
-
     text_labels.update({"header", "footer"})
     abandon_labels.discard("header")
     abandon_labels.discard("footer")
-
-    # Keep image-only decorative zones out of OCR by default.
-    # If you want logo OCR too, move these to text_labels.
     skip_labels.update({"header_image", "footer_image"})
     abandon_labels.discard("header_image")
     abandon_labels.discard("footer_image")
-
     label_task_mapping["text"] = sorted(text_labels)
     label_task_mapping["abandon"] = sorted(abandon_labels)
     label_task_mapping["skip"] = sorted(skip_labels)
@@ -142,8 +126,6 @@ def get_parser():
         _start_vllm_if_needed()
         _configure_glmocr_to_keep_header_footer()
         from glmocr import GlmOcr
-
-        # Force local self-hosted pipeline and avoid MaaS fallback.
         _parser = GlmOcr(mode="selfhosted")
     return _parser
 
@@ -153,8 +135,6 @@ def _extract_md(result: Any) -> str:
         return ""
 
     def _close_unbalanced_tables(chunk: str) -> str:
-        # Some OCR chunks end mid-table near page boundaries, which can cause
-        # subsequent content (including page separators) to render inside a cell.
         opens = len(re.findall(r"<table\b", chunk, flags=re.IGNORECASE))
         closes = len(re.findall(r"</table>", chunk, flags=re.IGNORECASE))
         missing = max(0, opens - closes)
@@ -173,11 +153,135 @@ def _extract_md(result: Any) -> str:
     return _close_unbalanced_tables(str(md).strip()) if md else ""
 
 
+def _get_cell_text(cell) -> str:
+    """Extract clean text from a BeautifulSoup table cell."""
+    # Replace <br> tags with space before getting text
+    for br in cell.find_all("br"):
+        br.replace_with(" ")
+    text = cell.get_text(separator=" ", strip=True)
+    # Collapse multiple spaces
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+
+def _html_table_to_markdown(table_html: str) -> str:
+    """
+    Convert a single HTML table to a clean markdown table.
+    - Removes empty columns that only contain whitespace
+    - Handles colspan/rowspan by ignoring them (takes first value)
+    - Works for both regular tables and bordered tables
+    """
+    try:
+        from bs4 import BeautifulSoup
+        soup = BeautifulSoup(table_html, "html.parser")
+        table = soup.find("table")
+        if not table:
+            return table_html
+    except ImportError:
+        # Fallback: regex-based if bs4 not available
+        return _html_table_to_markdown_regex(table_html)
+
+    # Collect all rows from thead + tbody + direct tr
+    all_rows = []
+    for tr in table.find_all("tr"):
+        cells = []
+        for cell in tr.find_all(["td", "th"]):
+            cells.append(_get_cell_text(cell))
+        if cells:
+            all_rows.append(cells)
+
+    if not all_rows:
+        return table_html
+
+    # Normalize all rows to same width
+    max_cols = max(len(r) for r in all_rows)
+    all_rows = [r + [""] * (max_cols - len(r)) for r in all_rows]
+
+    # Remove columns that are empty in ALL rows
+    non_empty_cols = []
+    for col_idx in range(max_cols):
+        if any(all_rows[row_idx][col_idx] for row_idx in range(len(all_rows))):
+            non_empty_cols.append(col_idx)
+
+    all_rows = [[row[i] for i in non_empty_cols] for row in all_rows]
+
+    if not all_rows:
+        return table_html
+
+    num_cols = len(all_rows[0])
+    md_lines = []
+
+    # First row as header
+    md_lines.append("| " + " | ".join(all_rows[0]) + " |")
+    md_lines.append("| " + " | ".join(["---"] * num_cols) + " |")
+
+    # Remaining rows
+    for row in all_rows[1:]:
+        # Pad or trim to match header width
+        padded = (row + [""] * num_cols)[:num_cols]
+        md_lines.append("| " + " | ".join(padded) + " |")
+
+    return "\n".join(md_lines)
+
+
+def _html_table_to_markdown_regex(table_html: str) -> str:
+    """Regex-based fallback for HTML table conversion when bs4 is unavailable."""
+    row_re = re.compile(r"<tr[^>]*>(.*?)</tr>", re.IGNORECASE | re.DOTALL)
+    cell_re = re.compile(r"<t[dh][^>]*>(.*?)</t[dh]>", re.IGNORECASE | re.DOTALL)
+    tag_re = re.compile(r"<[^>]+>")
+
+    def clean_cell(c: str) -> str:
+        c = re.sub(r"<br\s*/?>", " ", c, flags=re.IGNORECASE)
+        c = tag_re.sub("", c)
+        c = c.replace("&nbsp;", " ").replace("&amp;", "&").replace("&gt;", ">").replace("&lt;", "<")
+        return re.sub(r"\s+", " ", c).strip()
+
+    rows = []
+    for tr in row_re.findall(table_html):
+        cells = [clean_cell(c) for c in cell_re.findall(tr)]
+        if cells:
+            rows.append(cells)
+
+    if not rows:
+        return table_html
+
+    max_cols = max(len(r) for r in rows)
+    rows = [r + [""] * (max_cols - len(r)) for r in rows]
+
+    # Remove empty columns
+    non_empty_cols = [i for i in range(max_cols) if any(rows[j][i] for j in range(len(rows)))]
+    rows = [[r[i] for i in non_empty_cols] for r in rows]
+
+    if not rows:
+        return table_html
+
+    num_cols = len(rows[0])
+    md = ["| " + " | ".join(rows[0]) + " |", "| " + " | ".join(["---"] * num_cols) + " |"]
+    for row in rows[1:]:
+        padded = (row + [""] * num_cols)[:num_cols]
+        md.append("| " + " | ".join(padded) + " |")
+
+    return "\n".join(md)
+
+
+def _convert_all_html_tables(md: str) -> str:
+    """Find and replace all HTML tables in the markdown with clean markdown tables."""
+    table_re = re.compile(r"<table[\s\S]*?</table>", re.IGNORECASE)
+
+    def replace_table(match: re.Match) -> str:
+        return _html_table_to_markdown(match.group(0))
+
+    return table_re.sub(replace_table, md)
+
+
 def _clean_markdown(md: str) -> str:
     if not md:
         return md
 
-    # Drop frequently repeated bank footer/header lines.
+    # Convert all HTML tables to clean markdown tables first
+    md = _convert_all_html_tables(md)
+
+    # Drop frequently repeated bank footer/header lines
     drop_line_patterns = [
         r"Call 1-800-937-2000 for 24-hour Bank-by-Phone services.*",
         r"Bank Deposits FDIC Insured \| TD Bank, N\.A\. \| Equal Housing Lender.*",
@@ -194,8 +298,7 @@ def _clean_markdown(md: str) -> str:
         cleaned_lines.append(line)
     md = "\n".join(cleaned_lines)
 
-    # Remove long legal notice section that repeats and is usually not needed
-    # for transaction extraction.
+    # Remove long legal notice sections
     legal_start_markers = [
         "## FOR CONSUMER ACCOUNTS ONLY",
         "## FOR CONSUMER LOAN ACCOUNTS ONLY",
@@ -217,9 +320,13 @@ def _clean_markdown(md: str) -> str:
             kept.append(line)
     md = "\n".join(kept)
 
-    # Remove duplicate separators and excessive blank lines.
+    # Remove duplicate separators and excessive blank lines
     md = re.sub(r"(?:\n\s*){3,}", "\n\n", md)
-    md = re.sub(r"(?:\n\s*---page-separator---\s*\n){2,}", "\n\n---page-separator---\n\n", md)
+    md = re.sub(
+        r"(?:\n\s*---page-separator---\s*\n){2,}",
+        "\n\n---page-separator---\n\n",
+        md,
+    )
     return md.strip()
 
 
@@ -233,7 +340,6 @@ def run_ocr(file_obj):
         is_pdf = path.lower().endswith(".pdf")
         if is_pdf:
             page_images = _render_pdf_pages_to_images(path)
-            # Parse one rendered page per request item so we can preserve page boundaries.
             result = parser.parse(page_images)
         else:
             result = parser.parse(path)
@@ -241,13 +347,16 @@ def run_ocr(file_obj):
         return md
     except Exception as e:
         import traceback
-
         log.exception("run_ocr failed: %s", e)
         return f"Error: {e}\n\n{traceback.format_exc()}"
     finally:
         for p in page_images:
             try:
-                if isinstance(p, str) and p.endswith(".png") and "glmocr_page_" in os.path.basename(p):
+                if (
+                    isinstance(p, str)
+                    and p.endswith(".png")
+                    and "glmocr_page_" in os.path.basename(p)
+                ):
                     os.unlink(p)
             except Exception:
                 pass
@@ -271,4 +380,4 @@ def build_demo():
 
 
 if __name__ == "__main__":
-    build_demo().launch()
+    build_demo().launch()