Update app.py

app.py

@@ -1,19 +1,21 @@
 import gradio as gr
 import yaml
 import re
-import json
 import os
-import torch
 
+# Paths from installed glmocr package
 import glmocr
-GLMOCR_BASE    = os.path.dirname(glmocr.__file__)
-config_path    = os.path.join(GLMOCR_BASE, "config.yaml")
+GLMOCR_BASE = os.path.dirname(glmocr.__file__)
+config_path = os.path.join(GLMOCR_BASE, "config.yaml")
 formatter_path = os.path.join(GLMOCR_BASE, "postprocess", "result_formatter.py")
 
-# ── STEP 1: Fix config.yaml ───────────────────────────────────
+# ── STEP 1: Load and fix config (MaaS + keep header/footer) ─────────────────
 with open(config_path, "r") as f:
     config = yaml.safe_load(f)
 
+config["pipeline"]["maas"]["enabled"] = True
+config["pipeline"]["maas"]["api_key"] = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
+
 config["pipeline"]["result_formatter"]["abandon"] = [
     "number", "footnote", "aside_text",
     "reference", "footer_image", "header_image",
@@ -23,17 +25,15 @@ config["pipeline"]["enable_layout"] = True
 with open(config_path, "w") as f:
     yaml.dump(config, f, default_flow_style=False, sort_keys=False)
 
-print("✅ config.yaml fixed")
+print("✅ config.yaml fixed (MaaS enabled, header/footer kept)")
 
-# ── STEP 2: Fix result_formatter.py ──────────────────────────
+# ── STEP 2: Fix result_formatter.py ────────────────────────────────────────
 with open(formatter_path, "r") as f:
     source = f.read()
 
 labels_to_remove = [
-    '"header"', "'header'",
-    '"footer"', "'footer'",
-    '"doc_header"', "'doc_header'",
-    '"doc_footer"', "'doc_footer'"
+    '"header"', "'header'", '"footer"', "'footer'",
+    '"doc_header"', "'doc_header'", '"doc_footer"', "'doc_footer'"
 ]
 for label in labels_to_remove:
     source = re.sub(r',\s*' + re.escape(label), '', source)
@@ -45,154 +45,109 @@ with open(formatter_path, "w") as f:
 
 print("✅ result_formatter.py fixed")
 
-# ── STEP 3: Lazy model load ───────────────────────────────────
-processor = None
-model     = None
-ABANDON   = set(config["pipeline"]["result_formatter"]["abandon"])
-
-def load_model():
-    global processor, model
-    if model is not None:
-        return
-    from transformers import AutoProcessor, GlmOcrForConditionalGeneration
-    print("Loading model...")
-    processor = AutoProcessor.from_pretrained("zai-org/GLM-OCR")
-    model = GlmOcrForConditionalGeneration.from_pretrained(
-        "zai-org/GLM-OCR",
-        torch_dtype=torch.bfloat16,
-        device_map="auto",
-    )
-    print(f"✅ Model ready on {next(model.parameters()).device}")
-
-# ── STEP 4: OCR one image ─────────────────────────────────────
-def ocr_image(img_path):
-    messages = [
-        {"role": "user", "content": [
-            {"type": "image", "url": img_path},
-            {"type": "text",  "text": "Document Parsing:"}
-        ]}
-    ]
-    inputs = processor.apply_chat_template(
-        messages, tokenize=True, add_generation_prompt=True,
-        return_dict=True, return_tensors="pt"
-    ).to(model.device)
-    inputs.pop("token_type_ids", None)
-
-    with torch.no_grad():
-        output_ids = model.generate(**inputs, max_new_tokens=2048)
-
-    raw = processor.decode(
-        output_ids[0][inputs["input_ids"].shape[1]:],
-        skip_special_tokens=False
-    )
-
-    for token in ["<|assistant|>", "<|user|>", "<|system|>",
-                  "<|endoftext|>", "</s>", "<s>"]:
-        raw = raw.replace(token, "")
-    raw = raw.strip()
+ABANDON = set(config["pipeline"]["result_formatter"]["abandon"])
 
-    json_match = re.search(r'\[.*\]', raw, re.DOTALL)
-    if json_match:
-        try:
-            regions = json.loads(json_match.group())
-            return regions, raw, "json"
-        except:
-            pass
-
-    return [], raw, "raw"
-
-# ── STEP 5: Main OCR function ─────────────────────────────────
+# ── STEP 3: OCR via Zhipu MaaS (no local model) ─────────────────────────────
 def run_ocr(uploaded_file):
     if uploaded_file is None:
         return "Please upload a file.", "No regions detected."
 
+    api_key = "4570c28bdea5493c9efae9dae68edc66.sGbA9DLlcX1GlvqV"
+    if not api_key:
+        return (
+            "Error: ZHIPU_API_KEY is not set. Add it in Space Settings → Variables and secrets.",
+            "Failed."
+        )
+
     try:
-        load_model()
+        from glmocr import parse
 
         path = uploaded_file.name if hasattr(uploaded_file, "name") else str(uploaded_file)
 
+        # PDF → list of image paths (SDK expects images per page)
         if path.lower().endswith(".pdf"):
             import fitz
             doc = fitz.open(path)
             page_images = []
             for i in range(len(doc)):
-                pix      = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
-                img_path = f"/tmp/page_{i}.png"
+                pix = doc[i].get_pixmap(matrix=fitz.Matrix(1.5, 1.5), alpha=False)
+                img_path = f"/tmp/maas_page_{i}.png"
                 pix.save(img_path)
                 page_images.append(img_path)
             doc.close()
+            # parse() with list = one document, multiple pages
+            result = parse(page_images)
         else:
-            page_images = [path]
+            result = parse(path)
+
+        # result.json_result: list of pages, each page = list of regions
+        # Region: {"label": "header"|"footer"|"text"|..., "content": "...", ...}
+        pages_data = result.json_result if hasattr(result, "json_result") else []
+        if not isinstance(pages_data, list):
+            pages_data = [pages_data] if pages_data else []
 
         total_headers = 0
         total_footers = 0
-        all_markdown  = []
-        all_summary   = []
-
-        for page_num, img_path in enumerate(page_images):
-            regions, raw, output_type = ocr_image(img_path)
-
-            page_md      = []
-            page_summary = [f"── PAGE {page_num + 1} of {len(page_images)} ──"]
-
-            if output_type == "raw":
-                page_md.append(raw)
-                page_summary.append("[raw HTML/text output]")
-                page_summary.append(f"Content length: {len(raw)} chars")
-            else:
-                for region in regions:
-                    label   = region.get("label", "text")
-                    content = str(region.get("content", ""))
-
-                    if label in ABANDON:
-                        continue
-
-                    if label == "header":
-                        total_headers += 1
-                        page_summary.append(f"🔵 HEADER: {content[:100]}")
-                        page_md.append(f"<!-- HEADER -->\n{content}")
-                    elif label == "footer":
-                        total_footers += 1
-                        page_summary.append(f"🟢 FOOTER: {content[:100]}")
-                        page_md.append(f"<!-- FOOTER -->\n{content}")
-                    else:
-                        page_summary.append(f"[{label}]: {content[:100]}")
-                        page_md.append(content)
+        all_markdown = []
+        all_summary = []
+
+        for page_num, page_regions in enumerate(pages_data):
+            if not isinstance(page_regions, list):
+                page_regions = [page_regions] if page_regions else []
+            page_md = []
+            page_summary = [f"── PAGE {page_num + 1} of {len(pages_data)} ──"]
+
+            for region in page_regions:
+                if not isinstance(region, dict):
+                    continue
+                label = region.get("label", "text")
+                content = str(region.get("content", ""))
+
+                if label in ABANDON:
+                    continue
+
+                if label == "header":
+                    total_headers += 1
+                    page_summary.append(f"🔵 HEADER: {content[:100]}")
+                    page_md.append(f"<!-- HEADER -->\n{content}")
+                elif label == "footer":
+                    total_footers += 1
+                    page_summary.append(f"🟢 FOOTER: {content[:100]}")
+                    page_md.append(f"<!-- FOOTER -->\n{content}")
+                else:
+                    page_summary.append(f"[{label}]: {content[:100]}")
+                    page_md.append(content)
 
             all_markdown.append("\n\n".join(page_md))
             all_summary.extend(page_summary)
             all_summary.append("")
 
         summary = (
-            f"Total pages   : {len(page_images)}\n"
+            f"Total pages   : {len(pages_data)}\n"
             f"Headers found : {total_headers}\n"
             f"Footers found : {total_footers}\n"
             f"{'─'*40}\n"
             + "\n".join(all_summary)
         )
-
-        markdown = "\n\n---\n\n".join(all_markdown)
+        markdown = "\n\n---\n\n".join(all_markdown) if all_markdown else "(No content)"
         return markdown, summary
 
     except Exception as e:
         import traceback
         return f"Error: {str(e)}\n\n{traceback.format_exc()}", "Failed."
 
-# ── STEP 6: Gradio UI ─────────────────────────────────────────
-with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
+# ── STEP 4: Gradio UI ──────────────────────────────────────────────────────
+with gr.Blocks(title="GLM-OCR — MaaS (Header & Footer)") as demo:
     gr.Markdown("""
-    # 🔍 GLM-OCR — Header & Footer Fix
-    Upload PDF or image. Headers 🔵 and Footers 🟢 now appear in output.
-    Multi-page PDFs fully supported.
-    > First request takes ~2 min to load model. After that it is fast.
+    # 🔍 GLM-OCR — Zhipu MaaS
+    Upload PDF or image. Headers 🔵 and Footers 🟢 from the **full pipeline** (layout + OCR).
+    Uses your Zhipu API key (set in Space secrets).
     """)
-
     file_input = gr.File(
         label="Upload PDF or Image",
         file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".bmp"]
     )
     run_btn = gr.Button("▶  Run OCR", variant="primary", size="lg")
-
     with gr.Row():
         with gr.Column():
             gr.Markdown("### 📄 Markdown Output")
@@ -200,8 +155,6 @@ with gr.Blocks(title="GLM-OCR — Header & Footer Fix") as demo:
         with gr.Column():
             gr.Markdown("### 🗂️ Detected Regions")
             regions_out = gr.Textbox(lines=30, label="")
+    run_btn.click(fn=run_ocr, inputs=file_input, outputs=[markdown_out, regions_out])
 
-    run_btn.click(fn=run_ocr, inputs=file_input,
-                  outputs=[markdown_out, regions_out])
-
-demo.launch()
+demo.launch()