Update app.py

app.py

@@ -202,6 +202,7 @@ def _strip_code_fences(text: str) -> str:
 
 def _postprocess_markdown(text: str) -> str:
     text = _strip_code_fences(text)
+    text = _html_to_markdown_tables(text)
     text = _normalize_tables(text)
     return text.strip()
 
@@ -215,11 +216,14 @@ def _build_prompt(strict_table_mode: bool = False) -> str:
         "3) Never drop right-most numeric columns (amount/balance).\n"
         "4) Keep page content order exactly.\n"
         "5) Do not summarize.\n"
+        "6) Output markdown only.\n"
+        "7) Do not output HTML tags (<table>, <tr>, <td>, <thead>, <tbody>, etc.).\n"
+        "8) Do not invent sample/demo rows.\n"
     )
     if strict_table_mode:
         base += (
-            "6) If a row starts with a date (MM/DD), output it as a table row and include trailing amount.\n"
-            "7) Prefer explicit table rows over plain text for statement activity.\n"
+            "9) If a row starts with a date (MM/DD), output it as a table row and include trailing amount.\n"
+            "10) Prefer explicit table rows over plain text for statement activity.\n"
         )
     return base
 
@@ -232,6 +236,69 @@ def _count_dated_rows(text: str) -> int:
     return len(re.findall(r"(?m)^\s*\d{2}/\d{2}\b", text))
 
 
+def _html_to_markdown_tables(text: str) -> str:
+    """
+    Convert basic HTML table blocks into pipe markdown tables.
+    This keeps downstream parsing deterministic even if model emits HTML.
+    """
+    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)
+    br_re = re.compile(r"<br\s*/?>", re.IGNORECASE)
+    tag_re = re.compile(r"<[^>]+>")
+
+    def _clean_cell(s: str) -> str:
+        s = br_re.sub(" ", s)
+        s = tag_re.sub(" ", s)
+        s = s.replace("&nbsp;", " ")
+        s = s.replace("&amp;", "&")
+        s = re.sub(r"\s+", " ", s).strip()
+        return s
+
+    def _convert_one(table_html: str) -> str:
+        rows = []
+        for row_html in row_re.findall(table_html):
+            cells = [_clean_cell(c) for c in cell_re.findall(row_html)]
+            cells = [c for c in cells if c != ""]
+            if cells:
+                rows.append(cells)
+        if not rows:
+            return table_html
+
+        width = max(len(r) for r in rows)
+        norm_rows = [r + [""] * (width - len(r)) for r in rows]
+        header = norm_rows[0]
+        out = [
+            "| " + " | ".join(header) + " |",
+            "| " + " | ".join(["---"] * width) + " |",
+        ]
+        for r in norm_rows[1:]:
+            out.append("| " + " | ".join(r) + " |")
+        return "\n".join(out)
+
+    return table_re.sub(lambda m: _convert_one(m.group(0)), text)
+
+
+def _is_hallucinated_table_block(text: str) -> bool:
+    """
+    Detect obvious synthetic outputs like repeated '100.00' rows or
+    placeholder transaction lists not present in statement pages.
+    """
+    low = text.lower()
+    if "<table" not in low and "| transaction | amount |" not in low:
+        return False
+
+    repeated_100 = len(re.findall(r"\b100\.00\b", text))
+    credit_card_rows = len(re.findall(r"\bcredit card\b", low))
+    bank_transfer_rows = len(re.findall(r"\bbank transfer\b", low))
+
+    if repeated_100 >= 18 and (credit_card_rows + bank_transfer_rows) >= 10:
+        return True
+    if re.search(r"<!--\s*row\s+\d+\s*-->", low):
+        return True
+    return False
+
+
 def _looks_amount_missing(text: str) -> bool:
     """
     Heuristic: many dated activity rows but very low amount density.
@@ -257,7 +324,11 @@ def _looks_amount_missing(text: str) -> bool:
     return dated >= 12 and amts <= max(3, dated // 10)
 
 
-def _infer_image(image_path: str, strict_table_mode: bool = False) -> str:
+def _infer_image(
+    image_path: str,
+    strict_table_mode: bool = False,
+    conservative_mode: bool = False,
+) -> str:
     """Run fine-tuned model on a single image file and return markdown."""
     import torch
     from PIL import Image
@@ -272,12 +343,19 @@ def _infer_image(image_path: str, strict_table_mode: bool = False) -> str:
     try:
         img.save(resized_path, "PNG")
 
+        prompt = _build_prompt(strict_table_mode)
+        if conservative_mode:
+            prompt += (
+                "\n11) If unsure, keep original text line-by-line; do not fabricate values.\n"
+                "12) Never generate template/example/sample transaction rows.\n"
+            )
+
         messages = [
             {
                 "role": "user",
                 "content": [
                     {"type": "image", "url": resized_path},
-                    {"type": "text", "text": _build_prompt(strict_table_mode)},
+                    {"type": "text", "text": prompt},
                 ],
             }
         ]
@@ -374,7 +452,9 @@ def run_ocr(uploaded_file):
         all_pages: List[str] = []
 
         if not is_pdf:
-            page_md = _infer_image(path, strict_table_mode=True)
+            page_md = _infer_image(path, strict_table_mode=True, conservative_mode=True)
+            if _is_hallucinated_table_block(page_md):
+                page_md = _infer_image(path, strict_table_mode=True, conservative_mode=True)
             if page_md:
                 all_pages.append(page_md)
         else:
@@ -387,6 +467,20 @@ def run_ocr(uploaded_file):
                 temp_paths.append(img_path)
                 page_md = _infer_image(img_path, strict_table_mode=False)
 
+                # Pass 1.5: recover from synthetic/hallucinated table outputs
+                if _is_hallucinated_table_block(page_md):
+                    log.info(
+                        "Page %d flagged as hallucinated-table; retrying conservative prompt.",
+                        page_idx + 1,
+                    )
+                    page_md_retry = _infer_image(
+                        img_path,
+                        strict_table_mode=True,
+                        conservative_mode=True,
+                    )
+                    if not _is_hallucinated_table_block(page_md_retry):
+                        page_md = page_md_retry
+
                 # Pass 2: retry for amount-missing pages
                 if _looks_amount_missing(page_md):
                     log.info(
@@ -395,7 +489,11 @@ def run_ocr(uploaded_file):
                     )
                     retry_img = _render_pdf_page_to_image(path, page_idx, RETRY_PROFILE)
                     temp_paths.append(retry_img)
-                    retry_md = _infer_image(retry_img, strict_table_mode=True)
+                    retry_md = _infer_image(
+                        retry_img,
+                        strict_table_mode=True,
+                        conservative_mode=True,
+                    )
 
                     # Keep better output by amount density.
                     if _count_amounts(retry_md) > _count_amounts(page_md):