Update app.py

app.py

@@ -775,6 +775,52 @@ def invoice_text_to_json(
     final_json = _prune_empty_items(final_json)
     return final_json
 
+from typing import Optional
+
+# ----- replace old run_ocr with unified dispatcher -----
+def run_pipeline(file: Optional[gr.File], raw_txt: Optional[str]) -> str:
+    """
+    Orchestrates two intake lanes:
+    1) If raw_txt is provided (non-empty), skip OCR → directly map to schema.
+    2) Else, run OCR on the uploaded file and map to schema.
+    """
+    raw_txt = (raw_txt or "").strip()
+
+    # Lane A: Raw text → JSON
+    if raw_txt:
+        try:
+            result_json = invoice_text_to_json(raw_txt)
+            return json.dumps(result_json, indent=2, ensure_ascii=False)
+        except Exception as e:
+            return f"Error while converting pasted text to JSON schema: {e}"
+
+    # Lane B: File → OCR → JSON
+    if not file:
+        return "No input received. Upload an image/PDF or paste raw text."
+
+    try:
+        name = (file.name or "").lower()
+
+        # Load as DocumentFile (handles PNG/JPG/PDF)
+        if name.endswith(".pdf"):
+            doc = DocumentFile.from_pdf(file=file.name)
+        else:
+            doc = DocumentFile.from_images([file.name])
+
+        # Inference
+        result = MODEL(doc)
+        exported = result.export()
+        text = _collect_text_from_export(exported)
+        if not text:
+            return "No text detected by OCR."
+
+        result_json = invoice_text_to_json(text)
+        return json.dumps(result_json, indent=2, ensure_ascii=False)
+
+    except Exception as e:
+        return f"OCR pipeline error: {e}"
+        
+
 # ---------- Gradio UI ----------
 TITLE = "docTR OCR — Text Extractor"
 DESC = (
@@ -785,28 +831,43 @@ DESC = (
 with gr.Blocks(theme="soft", title=TITLE) as demo:
     gr.Markdown(f"# {TITLE}\n{DESC}")
 
-    with gr.Row():
-        inp = gr.File(label="Upload image/PDF", file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff", ".pdf"])
-    out = gr.Code(label="Extracted JSON", language="json")
-
+    with gr.TabbedInterface(
+        [
+            gr.TabItem("Upload File"),
+            gr.TabItem("Paste Raw Text")
+        ],
+        tab_names=["Upload File", "Paste Raw Text"]
+    ) as tabs:
+        with tabs.children[0]:
+            inp = gr.File(
+                label="Upload image/PDF",
+                file_types=[".png", ".jpg", ".jpeg", ".tif", ".tiff", .".pdf"]
+            )
+            # hidden paired textbox to keep a single click path
+            raw_txt_hidden = gr.Textbox(visible=False)
+
+        with tabs.children[1]:
+            # expose text lane; hide file lane counterpart
+            raw_txt = gr.Textbox(
+                label="Paste raw invoice text (we’ll map directly to JSON schema)",
+                lines=18,
+                placeholder="Paste the OCR’d/plain text of the invoice here…"
+            )
+            file_hidden = gr.File(visible=False)
 
-    run_btn = gr.Button("Run OCR", variant="primary")
-    run_btn.click(fn=run_ocr, inputs=inp, outputs=out)
+    out = gr.Code(label="Extracted JSON", language="json")
+    run_btn = gr.Button("Generate JSON", variant="primary")
 
-    gr.Examples(
-        examples=[
-            # You can drop a couple of public sample URLs here if desired,
-            # but Spaces won't auto-download without code. Leave empty by default.
-        ],
-        inputs=inp,
+    # One button → unified function; we pass both lanes (visible/hidden)
+    run_btn.click(
+        fn=run_pipeline,
+        inputs=[inp, raw_txt],
         outputs=out,
-        cache_examples=False,
-        label="(Optional) Examples"
     )
 
     gr.Markdown(
-        "Tip: For multi-page PDFs, the output shows a **PAGE BREAK** separator between pages.\n"
-        "For production pipelines, capture this output and route it to your parsing/LLM layer."
+        "ℹ️ **Usage:** Prefer the *Paste Raw Text* tab when you already have text and only need schema mapping. "
+        "If both file and text are provided, we’ll **prioritize the pasted text** by design to collapse cycle time."
     )
 
 if __name__ == "__main__":