Update app.py

app.py

@@ -87,57 +87,140 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
+def _build_table(rows: List[str]) -> str:
+    parsed = []
+
+    for row in rows:
+        parts = re.split(r"\s{2,}", row.strip())
+        if len(parts) >= 3:
+            date = parts[0]
+            amount = parts[-1] if re.search(r"\d+\.\d{2}", parts[-1]) else ""
+            desc = " ".join(parts[1:-1]) if amount else " ".join(parts[1:])
+            parsed.append((date, desc, amount))
+
+    if not parsed:
+        return "\n".join(rows)
+
+    table = [
+        "| POSTING DATE | DESCRIPTION | AMOUNT |",
+        "| :--- | :--- | ---: |",
+    ]
+    for date, desc, amount in parsed:
+        table.append(f"| {date} | {desc} | {amount} |")
+    return "\n".join(table)
+
+
+def _normalize_transactions(text: str) -> str:
+    """
+    Convert loose transaction lines into structured markdown tables.
+    """
+    lines = text.split("\n")
+    result: List[str] = []
+    table_buffer: List[str] = []
+    in_section = False
+
+    for line in lines:
+        if re.search(r"POSTING DATE.*DESCRIPTION.*AMOUNT", line, re.IGNORECASE):
+            in_section = True
+            table_buffer = []
+            continue
+
+        if in_section and (line.strip() == "" or line.strip().startswith("Subtotal")):
+            if table_buffer:
+                result.append(_build_table(table_buffer))
+                table_buffer = []
+            in_section = False
+            result.append(line)
+            continue
+
+        if in_section:
+            table_buffer.append(line)
+        else:
+            result.append(line)
+
+    if table_buffer:
+        result.append(_build_table(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:
+        return 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 _clean_markdown(text: str) -> str:
-    """Post-process markdown to fix common formatting issues."""
-    
-    # Fix table formatting - ensure proper spacing
+    """Post-process markdown to fix table formatting and section structure."""
+    text = _html_to_markdown_tables(text)
+    text = _fix_missing_amounts(text)
+    text = _normalize_transactions(text)
+
     lines = text.split('\n')
     cleaned = []
     in_table = False
-    
-    for i, line in enumerate(lines):
-        # Detect table rows
+
+    for line in lines:
         if '|' in line and line.strip().startswith('|'):
             if not in_table:
-                # First table row - add blank line before if needed
                 if cleaned and cleaned[-1].strip():
                     cleaned.append('')
                 in_table = True
-            
-            # Clean up table row spacing
+
             line = re.sub(r'\s*\|\s*', ' | ', line)
             line = re.sub(r'\s+', ' ', line)
             cleaned.append(line.strip())
-            
         elif in_table and line.strip() == '':
-            # End of table - add blank line after
             in_table = False
             cleaned.append('')
-            
         else:
             in_table = False
-            # Remove excessive blank lines
             if line.strip() or (cleaned and cleaned[-1].strip()):
                 cleaned.append(line)
-    
-    # Join lines
+
     text = '\n'.join(cleaned)
-    
-    # Fix header spacing
     text = re.sub(r'(#{1,6})\s*([^\n]+)', r'\1 \2', text)
-    
-    # Fix list spacing
     text = re.sub(r'\n([•\-\*])\s+', r'\n\1 ', text)
-    
-    # Remove trailing whitespace from lines
     text = '\n'.join(line.rstrip() for line in text.split('\n'))
-    
-    # Normalize multiple blank lines to maximum 2
     text = re.sub(r'\n{4,}', '\n\n\n', text)
-    
-    # Fix common OCR artifacts in tables
-    text = re.sub(r'\|\s*:\s*---\s*\|\s*:\s*---\s*\|', '| :--- | :--- |', text)
-    
     return text.strip()
 
 
@@ -160,7 +243,17 @@ def _infer_image(image_path: str) -> str:
             "role": "user",
             "content": [
                 {"type": "image", "url": resized_path},
-                {"type": "text", "text": "Document Parsing:"},
+                {
+                    "type": "text",
+                    "text": (
+                        "Document Parsing to markdown.\n"
+                        "Rules:\n"
+                        "1) Preserve rows exactly in reading order.\n"
+                        "2) Keep transaction blocks as markdown tables.\n"
+                        "3) Do not invent rows or sample/template data.\n"
+                        "4) Keep right-most amount values."
+                    ),
+                },
             ],
         }]
 
@@ -173,7 +266,8 @@ def _infer_image(image_path: str) -> str:
         ).to(model.device)
         inputs.pop("token_type_ids", None)
 
-        torch.cuda.empty_cache()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
 
         with torch.no_grad():
             ids = model.generate(
@@ -305,12 +399,6 @@ def _create_gradio_demo():
         )
         run_btn.click(fn=run_ocr, inputs=file_in, outputs=out)
         
-        gr.Markdown("""
-        ### Tips
-        - Tables are formatted in markdown format
-        - Use the copy button to export the markdown
-        - Each page is separated by `---page-separator---`
-        """)
     return demo