Update app.py

app.py

@@ -87,285 +87,111 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
-def _build_table_from_entries(entries: List[dict]) -> str:
-    if not entries:
-        return ""
-    table = [
-        "| Date | Description | Amount |",
-        "|------|------------|--------|",
-    ]
-    for row in entries:
-        table.append(f"| {row['date']} | {row['description']} | {row['amount']} |")
-    return "\n".join(table)
+def _detect_headers(header_line: str) -> List[str]:
+    """Extract column headers from a header line by splitting on 2+ spaces."""
+    parts = re.split(r'\s{2,}', header_line.strip())
+    headers = [p.strip() for p in parts if p.strip()]
+    return headers if headers else ["POSTING DATE", "DESCRIPTION", "AMOUNT"]
 
 
-def _is_transaction_start(line: str) -> bool:
-    return bool(re.match(r"\d{2}/\d{2}", line.strip()))
-
-
-def _ends_with_amount(text: str) -> bool:
-    return bool(re.search(r"\d{1,3}(?:,\d{3})*\.\d{2}$", text.strip()))
-
-
-def _is_noise_line(line: str) -> bool:
-    s = line.strip().lower()
-    if not s:
-        return True
-    if s == "---page-separator---":
-        return True
-    if "call 1-800-937-2000" in s:
-        return True
-    if s in ("<table border=\"1\"><tr><td></td></tr></table>", "<table><tr><td></td></tr></table>"):
-        return True
-    return False
-
-
-def _merge_transaction_lines(rows: List[str]) -> List[str]:
-    """
-    Step 1/2: Merge multiline OCR rows into single transaction rows.
-    A new row begins when a line starts with MM/DD.
-    """
-    merged: List[str] = []
-    current = ""
-
-    for raw in rows:
-        line = " ".join(raw.strip().split())
-        if not line or _is_noise_line(line):
-            continue
-
-        if _is_transaction_start(line):
-            if current:
-                merged.append(current.strip())
-            current = line
-        else:
-            if current:
-                current += " " + line
-            else:
-                current = line
-
-    if current:
-        merged.append(current.strip())
-
-    return merged
-
-
-def _repair_orphan_amounts(rows: List[str]) -> List[str]:
+def _plaintext_rows_to_table(text: str) -> str:
     """
-    Step 3: Attach lines without ending amount to prior transaction line when possible.
+    Convert plain-text transaction rows (no pipe boundaries) into proper
+    markdown tables. Detects headers dynamically from the document.
+    Works for any column layout, not just Date/Description/Amount.
     """
-    fixed: List[str] = []
-    for line in rows:
-        if _ends_with_amount(line):
-            fixed.append(line)
-        elif fixed:
-            fixed[-1] = f"{fixed[-1]} {line}".strip()
-        else:
-            fixed.append(line)
-    return fixed
-
+    lines = text.split('\n')
+    result = []
+    i = 0
+    # Pattern to detect a header line (2+ words separated by 2+ spaces)
+    header_re = re.compile(
+        r'^[A-Z][A-Z\s]{3,}$',
+    )
+    # Pattern for a transaction row starting with MM/DD
+    txn_re = re.compile(r'^\d{2}/\d{2}\s+\S')
+    # Pattern to extract amount at end of line
+    amt_re = re.compile(r'^(.*?)\s+([\d,]+\.\d{2})$')
 
-def _normalize_transactions(text: str) -> str:
-    """
-    Convert loose transaction lines into structured markdown tables using:
-    MM/DD ... amount
-    """
-    lines = text.split("\n")
-    result: List[str] = []
-    table_buffer: List[str] = []
-    in_section = False
-    pattern = re.compile(r"(\d{2}/\d{2})\s+(.*?)\s+([\d,]+\.\d{2})$")
-
-    for line in lines:
-        if re.search(r"POSTING DATE.*DESCRIPTION.*AMOUNT", line, re.IGNORECASE):
-            in_section = True
-            table_buffer = []
-            continue
+    while i < len(lines):
+        line = lines[i]
+        stripped = line.strip()
 
-        if in_section and (line.strip() == "" or line.strip().startswith("Subtotal")):
-            if table_buffer:
-                clean_rows = []
-                merged_rows = _merge_transaction_lines(table_buffer)
-                repaired_rows = _repair_orphan_amounts(merged_rows)
-                for normalized in repaired_rows:
-                    match = pattern.search(normalized)
-                    if match:
-                        date, desc, amount = match.groups()
-                        clean_rows.append(
-                            {"date": date, "description": desc, "amount": amount}
-                        )
-                if clean_rows:
-                    result.append(_build_table_from_entries(clean_rows))
-                else:
-                    result.extend(table_buffer)
-                table_buffer = []
-            in_section = False
+        # Already a proper markdown table row — pass through unchanged
+        if stripped.startswith('|'):
             result.append(line)
+            i += 1
             continue
 
-        if in_section:
-            table_buffer.append(line)
-        else:
-            result.append(line)
-
-    if table_buffer:
-        clean_rows = []
-        merged_rows = _merge_transaction_lines(table_buffer)
-        repaired_rows = _repair_orphan_amounts(merged_rows)
-        for normalized in repaired_rows:
-            match = pattern.search(normalized)
-            if match:
-                date, desc, amount = match.groups()
-                clean_rows.append(
-                    {"date": date, "description": desc, "amount": amount}
-                )
-        if clean_rows:
-            result.append(_build_table_from_entries(clean_rows))
-        else:
-            result.extend(table_buffer)
-
-    return "\n".join(result)
-
-
-def _fix_missing_amounts(text: str) -> str:
-    lines = text.split("\n")
-    fixed: List[str] = []
-
-    for line in lines:
-        if re.search(r"\d{2}/\d{2}", line) and not re.search(r"\d+\.\d{2}", line):
-            match = re.search(r"(\d{1,3}(?:,\d{3})*\.\d{2})$", line)
-            if match:
-                line += f" {match.group(1)}"
-        fixed.append(line)
-
-    return "\n".join(fixed)
-
-
-def _html_to_markdown_tables(text: str) -> str:
-    try:
-        from bs4 import BeautifulSoup
-    except Exception:
-        # Fallback when bs4 is unavailable: regex-based conversion for simple tables.
-        table_re = re.compile(r"<table[\s\S]*?</table>", re.IGNORECASE)
-        row_re = re.compile(r"<tr[\s\S]*?</tr>", re.IGNORECASE)
-        cell_re = re.compile(r"<t[dh][^>]*>([\s\S]*?)</t[dh]>", re.IGNORECASE)
-        tag_re = re.compile(r"<[^>]+>")
-
-        def _clean_cell(cell: str) -> str:
-            cell = re.sub(r"<br\s*/?>", " ", cell, flags=re.IGNORECASE)
-            cell = tag_re.sub(" ", cell)
-            cell = cell.replace("&nbsp;", " ").replace("&amp;", "&")
-            cell = re.sub(r"\s+", " ", cell).strip()
-            return cell
-
-        def _convert(table_html: str) -> str:
-            rows = []
-            for tr in row_re.findall(table_html):
-                cols = [_clean_cell(c) for c in cell_re.findall(tr)]
-                cols = [c for c in cols if c != ""]
-                if cols:
-                    rows.append(cols)
-            if not rows:
-                return table_html
-            width = max(len(r) for r in rows)
-            norm = [r + [""] * (width - len(r)) for r in rows]
-            md = []
-            md.append("| " + " | ".join(norm[0]) + " |")
-            md.append("| " + " | ".join(["---"] * width) + " |")
-            for r in norm[1:]:
-                md.append("| " + " | ".join(r) + " |")
-            return "\n".join(md)
-
-        return table_re.sub(lambda m: _convert(m.group(0)), text)
-
-    soup = BeautifulSoup(text, "html.parser")
-
-    for table in soup.find_all("table"):
-        rows = []
-        for tr in table.find_all("tr"):
-            cols = [td.get_text(strip=True) for td in tr.find_all(["td", "th"])]
-            if cols:
-                rows.append(cols)
-
-        if rows:
-            md = []
-            header = rows[0]
-            md.append("| " + " | ".join(header) + " |")
-            md.append("| " + " | ".join(["---"] * len(header)) + " |")
-
-            for row in rows[1:]:
-                padded = row + [""] * (len(header) - len(row))
-                md.append("| " + " | ".join(padded[: len(header)]) + " |")
-
-            table.replace_with("\n".join(md))
-
-    return str(soup)
-
-
-def _remove_prompt_echo_artifacts(text: str) -> str:
-    """
-    Remove model prompt-echo garbage blocks such as repeated:
-    'N) Keep left-most amount values.' / 'N) Keep right-most amount values.'
-    """
-    lines = text.split("\n")
-    cleaned: List[str] = []
-    keep_line_re = re.compile(
-        r"^\s*\d+\)\s*Keep\s+(left-most|right-most)\s+amount\s+values\.?\s*$",
-        re.IGNORECASE,
-    )
-    repeat_run = 0
-
-    for line in lines:
-        if keep_line_re.match(line):
-            repeat_run += 1
+        # Detect start of a plain-text transaction block
+        if txn_re.match(stripped):
+            # Look back to find the most recent header line before this block
+            detected_headers = None
+            for j in range(len(result) - 1, max(len(result) - 5, -1), -1):
+                prev = result[j].strip()
+                if prev and not prev.startswith('|') and not prev.startswith('#'):
+                    candidate_headers = _detect_headers(prev)
+                    if len(candidate_headers) >= 2:
+                        detected_headers = candidate_headers
+                        # Remove the header line from result since we'll put it in table
+                        result.pop(j)
+                        break
+
+            if not detected_headers:
+                detected_headers = ["POSTING DATE", "DESCRIPTION", "AMOUNT"]
+
+            # Collect all transaction rows in this block
+            block = []
+            while i < len(lines):
+                l = lines[i].strip()
+                if not l:
+                    break
+                if l.startswith('|') or l.startswith('#'):
+                    break
+                if re.match(r'^Subtotal', l, re.IGNORECASE):
+                    break
+                if txn_re.match(l):
+                    block.append(l)
+                elif block:
+                    # Continuation line — merge into previous row
+                    block[-1] += ' ' + l
+                i += 1
+
+            if block:
+                # Build markdown table with detected headers
+                num_cols = len(detected_headers)
+                result.append('| ' + ' | '.join(detected_headers) + ' |')
+                result.append('| ' + ' | '.join(['---'] * num_cols) + ' |')
+
+                for row in block:
+                    if num_cols == 3:
+                        # Try to split into date / description / amount
+                        m = amt_re.match(row)
+                        if m:
+                            rest = m.group(1)
+                            amount = m.group(2)
+                            date_m = re.match(r'^(\d{2}/\d{2})\s+(.*)', rest)
+                            if date_m:
+                                result.append(
+                                    f'| {date_m.group(1)} | {date_m.group(2)} | {amount} |'
+                                )
+                                continue
+                        # No amount found — put full row in description, leave amount blank
+                        date_m = re.match(r'^(\d{2}/\d{2})\s+(.*)', row)
+                        if date_m:
+                            result.append(
+                                f'| {date_m.group(1)} | {date_m.group(2)} | |'
+                            )
+                        else:
+                            result.append(f'| | {row} | |')
+                    else:
+                        # For non-3-column tables just put full row as single cell
+                        result.append(f'| {row} |')
             continue
-        # Also drop explicit prompt/rules echoes if they appear as plain output.
-        if re.match(r"^\s*Document Parsing to markdown\.?\s*$", line, re.IGNORECASE):
-            continue
-        if re.match(r"^\s*Rules:\s*$", line, re.IGNORECASE):
-            continue
-        cleaned.append(line)
-
-    # If there was a long repeated run, ensure dangling separator spacing stays clean.
-    text = "\n".join(cleaned)
-    text = re.sub(r"\n{4,}", "\n\n\n", text)
-    return text
 
+        result.append(line)
+        i += 1
 
-def _clean_markdown(text: str) -> str:
-    """Post-process markdown to fix table formatting and section structure."""
-    text = _remove_prompt_echo_artifacts(text)
-    text = _html_to_markdown_tables(text)
-    text = _fix_missing_amounts(text)
-    text = _normalize_transactions(text)
-
-    lines = text.split('\n')
-    cleaned = []
-    in_table = False
-
-    for line in lines:
-        if '|' in line and line.strip().startswith('|'):
-            if not in_table:
-                if cleaned and cleaned[-1].strip():
-                    cleaned.append('')
-                in_table = True
-
-            line = re.sub(r'\s*\|\s*', ' | ', line)
-            line = re.sub(r'\s+', ' ', line)
-            cleaned.append(line.strip())
-        elif in_table and line.strip() == '':
-            in_table = False
-            cleaned.append('')
-        else:
-            in_table = False
-            if line.strip() or (cleaned and cleaned[-1].strip()):
-                cleaned.append(line)
-
-    text = '\n'.join(cleaned)
-    text = re.sub(r'(#{1,6})\s*([^\n]+)', r'\1 \2', text)
-    text = re.sub(r'\n([•\-\*])\s+', r'\n\1 ', text)
-    text = '\n'.join(line.rstrip() for line in text.split('\n'))
-    text = re.sub(r'\n{4,}', '\n\n\n', text)
-    return text.strip()
+    return '\n'.join(result)
 
 
 def _infer_image(image_path: str) -> str:
@@ -387,14 +213,7 @@ def _infer_image(image_path: str) -> str:
             "role": "user",
             "content": [
                 {"type": "image", "url": resized_path},
-                {
-                    "type": "text",
-                    "text": (
-                        "Extract document text as markdown only. "
-                        "Do not include instructions, examples, or template rows. "
-                        "Preserve transaction lines exactly."
-                    ),
-                },
+                {"type": "text", "text": "Document Parsing:"},
             ],
         }]
 
@@ -407,8 +226,7 @@ def _infer_image(image_path: str) -> str:
         ).to(model.device)
         inputs.pop("token_type_ids", None)
 
-        if torch.cuda.is_available():
-            torch.cuda.empty_cache()
+        torch.cuda.empty_cache()
 
         with torch.no_grad():
             ids = model.generate(
@@ -422,9 +240,9 @@ def _infer_image(image_path: str) -> str:
             ids[0][inputs["input_ids"].shape[1]:],
             skip_special_tokens=True,
         )
-        
-        # Clean up the markdown output
-        return _clean_markdown(result.strip())
+
+        result = _plaintext_rows_to_table(result)
+        return result.strip()
 
     finally:
         try:
@@ -528,18 +346,13 @@ def _create_gradio_demo():
     with gr.Blocks(title="GLM-OCR Fine-tuned") as demo:
         gr.Markdown("# GLM-OCR (Fine-tuned)")
         gr.Markdown("Upload a PDF or image to extract structured markdown content.")
-        
         file_in = gr.File(
             label="Upload PDF or image",
             file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"],
         )
         run_btn = gr.Button("Run OCR", variant="primary")
-        out = gr.Textbox(
-            lines=40, 
-            label="Output (markdown)"
-        )
+        out = gr.Textbox(lines=40, label="Output (markdown)")
         run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
-        
     return demo