Update app.py

app.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import re
 import tempfile
 import uuid
 from typing import List, Tuple
@@ -86,6 +87,60 @@ def _resize_for_inference(img):
     return img.resize(new_size, Image.LANCZOS)
 
 
+def _clean_markdown(text: str) -> str:
+    """Post-process markdown to fix common formatting issues."""
+    
+    # Fix table formatting - ensure proper spacing
+    lines = text.split('\n')
+    cleaned = []
+    in_table = False
+    
+    for i, line in enumerate(lines):
+        # Detect table rows
+        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()
+
+
 def _infer_image(image_path: str) -> str:
     """Run fine-tuned model on a single image file and return markdown string."""
     import torch
@@ -132,7 +187,9 @@ def _infer_image(image_path: str) -> str:
             ids[0][inputs["input_ids"].shape[1]:],
             skip_special_tokens=True,
         )
-        return result.strip()
+        
+        # Clean up the markdown output
+        return _clean_markdown(result.strip())
 
     finally:
         try:
@@ -235,13 +292,26 @@ 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)",
+            show_copy_button=True
+        )
         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