Update app.py

app.py

@@ -1,339 +1,306 @@
 # app.py
-# Invoice Extraction — Donut (public HF model) + Tesseract tables
-# Robust PDF handling:
-# 1) Try pdf2image with Poppler path detection (Fix A)
-# 2) If Poppler is missing, auto-fallback to PyMuPDF (no Poppler required)
+# Invoice -> JSON (Paste Text Only)
+# Pipeline:
+# 1) User pastes invoice text into a textbox (no OCR calls).
+# 2) Extract naive key:value candidates.
+# 3) Map candidates to your static schema headers with all-MiniLM-L6-v2.
+# 4) Ask MD2JSON-T5-small-V1 to emit STRICT JSON using your schema (nulls if missing).
+# If the model returns invalid JSON, we fall back to a schema-shaped object filled from mapped data (null where missing).
+
+import os
+import re
+import json
+from typing import List, Dict
 
-import os, io, re, json, shutil
-from typing import List
 import numpy as np
-import pandas as pd
-from PIL import Image, ImageOps, ImageFilter
-
 import streamlit as st
-
-# OCR (detection only) and PDF->image
-import pytesseract
-from pytesseract import Output
-from pdf2image import convert_from_bytes
-
-# HF Donut (public model)
 import torch
-from transformers import DonutProcessor, VisionEncoderDecoderModel
-
-# ------------------------------------------------------------------
-st.set_page_config(
-    page_title="Invoice Extraction — Donut (public) + Tesseract tables",
-    layout="wide"
-)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ----------------------------- Sidebar -----------------------------
-st.sidebar.header("Model (Hugging Face — public)")
-model_id = st.sidebar.text_input(
-    "HF model id",
-    value="naver-clova-ix/donut-base-finetuned-cord-v2",
-    help="Use a public model id; this one works without token."
-)
-task_prompt = st.sidebar.text_input(
-    "Task prompt (Donut)",
-    value="<s_cord-v2>",
-    help="Keep default for CORD-style invoices."
-)
-det_lang = st.sidebar.text_input("Tesseract language(s) — detection only", value="eng")
-show_boxes = st.sidebar.checkbox("Show word boxes (debug)", value=False)
-
-# ----------------------------- PDF loader (Fix A + fallback) -----------------------------
-def _find_poppler_path():
-    # Return a folder containing pdfinfo/pdftoppm if not on PATH
-    if shutil.which("pdfinfo") and shutil.which("pdftoppm"):
-        return None
-    for p in ["/usr/bin", "/usr/local/bin", "/usr/share/bin"]:
-        if os.path.exists(os.path.join(p, "pdfinfo")) and os.path.exists(os.path.join(p, "pdftoppm")):
-            return p
-    return None
-
-def _pages_via_pdf2image(file_bytes: bytes) -> List[Image.Image]:
-    poppler_path = _find_poppler_path()
-    if poppler_path:
-        return convert_from_bytes(file_bytes, dpi=300, poppler_path=poppler_path)
-    else:
-        return convert_from_bytes(file_bytes, dpi=300)
-
-def _pages_via_pymupdf(file_bytes: bytes) -> List[Image.Image]:
-    import fitz  # PyMuPDF
-    doc = fitz.open(stream=file_bytes, filetype="pdf")
-    pages = []
-    for page in doc:
-        # Use a mild upscale for better OCR if you want: matrix = fitz.Matrix(2, 2)
-        pix = page.get_pixmap()  # or: page.get_pixmap(matrix=matrix)
-        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-        pages.append(img)
-    return pages
+from transformers import pipeline
+from sentence_transformers import SentenceTransformer, util
+
+st.set_page_config(page_title="Invoice → JSON (Paste Text) · MiniLM + MD2JSON", layout="wide")
+st.title("Invoice → JSON (Paste Text Only)")
+
+# ----------------------------- Your schema (fixed) -----------------------------
+SCHEMA_JSON = {
+  "invoice_header": {
+    "car_number": None,
+    "shipment_number": None,
+    "shipping_point": None,
+    "currency": None,
+    "invoice_number": None,
+    "invoice_date": None,
+    "order_number": None,
+    "customer_order_number": None,
+    "our_order_number": None,
+    "sales_order_number": None,
+    "purchase_order_number": None,
+    "order_date": None,
+    "supplier_name": None,
+    "supplier_address": None,
+    "supplier_phone": None,
+    "supplier_email": None,
+    "supplier_tax_id": None,
+    "customer_name": None,
+    "customer_address": None,
+    "customer_phone": None,
+    "customer_email": None,
+    "customer_tax_id": None,
+    "ship_to_name": None,
+    "ship_to_address": None,
+    "bill_to_name": None,
+    "bill_to_address": None,
+    "remit_to_name": None,
+    "remit_to_address": None,
+    "tax_id": None,
+    "tax_registration_number": None,
+    "vat_number": None,
+    "payment_terms": None,
+    "payment_method": None,
+    "payment_reference": None,
+    "bank_account_number": None,
+    "iban": None,
+    "swift_code": None,
+    "total_before_tax": None,
+    "tax_amount": None,
+    "tax_rate": None,
+    "shipping_charges": None,
+    "discount": None,
+    "total_due": None,
+    "amount_paid": None,
+    "balance_due": None,
+    "due_date": None,
+    "invoice_status": None,
+    "reference_number": None,
+    "project_code": None,
+    "department": None,
+    "contact_person": None,
+    "notes": None,
+    "additional_info": None
+  },
+  "line_items": [
+    {
+      "quantity": None,
+      "units": None,
+      "description": None,
+      "footage": None,
+      "price": None,
+      "amount": None,
+      "notes": None
+    }
+  ]
+}
 
-def load_pages(file_bytes: bytes, name: str) -> List[Image.Image]:
-    name = (name or "").lower()
-    if name.endswith(".pdf"):
-        # Try Poppler route first
-        try:
-            return _pages_via_pdf2image(file_bytes)
-        except Exception:
-            # Fallback: PyMuPDF (no Poppler required)
-            return _pages_via_pymupdf(file_bytes)
-    return [Image.open(io.BytesIO(file_bytes)).convert("RGB")]
+STATIC_HEADERS: List[str] = list(SCHEMA_JSON["invoice_header"].keys())
 
-def preprocess_for_detection(img: Image.Image) -> Image.Image:
-    g = ImageOps.grayscale(img)
-    g = ImageOps.autocontrast(g)
-    g = g.filter(ImageFilter.UnsharpMask(radius=1, percent=150, threshold=3))
-    return g
+# ----------------------------- Sidebar controls -----------------------------
+st.sidebar.header("Settings")
+threshold = st.sidebar.slider("Semantic match threshold (cosine)", 0.0, 1.0, 0.60, 0.01)
+max_new_tokens = st.sidebar.slider("Max new tokens (MD2JSON)", 128, 2048, 512, 32)
+show_intermediates = st.sidebar.checkbox("Show intermediate mapping/hints", value=False)
 
-# ----------------------------- Donut loader -----------------------------
+# ----------------------------- Model loaders (cached) -----------------------------
 @st.cache_resource(show_spinner=True)
-def load_donut(_model_id: str):
-    processor = DonutProcessor.from_pretrained(_model_id)
-    model = VisionEncoderDecoderModel.from_pretrained(_model_id)
-    model.to(device).eval()
-    return processor, model
-
-def donut_infer(img: Image.Image, processor: DonutProcessor, model: VisionEncoderDecoderModel, prompt: str):
-    inputs = processor(images=img, text=prompt, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=1024, num_beams=1, early_stopping=True)
-    seq = processor.batch_decode(outputs, skip_special_tokens=True)[0]
-    parsed = None
-    try:
-        start = seq.find("{")
-        end = seq.rfind("}")
-        if start != -1 and end != -1 and end > start:
-            parsed = json.loads(seq[start:end+1])
-    except Exception:
-        parsed = None
-    return seq, parsed
-
-# ----------------------------- Key fields & tables -----------------------------
-CURRENCY = r"(?P<curr>USD|CAD|EUR|GBP|\$|C\$|€|£)?"
-MONEY = rf"{CURRENCY}\s?(?P<amt>\d{{1,3}}(?:[,]\d{{3}})*(?:[.]\d{{2}})?)"
-DATE = r"(?P<date>(?:\d{4}[-/]\d{1,2}[-/]\d{1,2})|(?:\d{1,2}[-/]\d{1,2}[-/]\d{2,4})|(?:[A-Za-z]{3,9}\s+\d{1,2},\s*\d{2,4}))"
-INV_PAT = r"(?:invoice\s*(?:no\.?|#|number)?\s*[:\-]?\s*(?P<inv>[A-Z0-9\-_/]{4,}))"
-PO_PAT  = r"(?:po\s*(?:no\.?|#|number)?\s*[:\-]?\s*(?P<po>[A-Z0-9\-_/]{3,}))"
-TOTAL_PAT = rf"(?:\b(total(?:\s*amount)?|amount\s*due|grand\s*total)\b.*?{MONEY})"
-SUBTOTAL_PAT = rf"(?:\bsub\s*total\b.*?{MONEY})"
-TAX_PAT = rf"(?:\b(tax|gst|vat|hst)\b.*?{MONEY})"
-
-def parse_fields_regex(fulltext: str):
-    t = re.sub(r"[ \t]+", " ", fulltext)
-    t = re.sub(r"\n{2,}", "\n", t)
-    out = {"invoice_number":None,"invoice_date":None,"po_number":None,"subtotal":None,"tax":None,"total":None,"currency":None}
-    m = re.search(INV_PAT, t, re.I);           out["invoice_number"] = m.group("inv") if m else None
-    m = re.search(PO_PAT, t, re.I);            out["po_number"] = m.group("po") if m else None
-    m = re.search(rf"(invoice\s*date[:\-\s]*){DATE}", t, re.I)
-    out["invoice_date"] = (m.group("date") if m else (re.search(DATE, t, re.I).group("date") if re.search(DATE, t, re.I) else None))
-    m = re.search(SUBTOTAL_PAT, t, re.I|re.S);
-    if m: out["subtotal"], out["currency"] = m.group("amt").replace(",",""), m.group("curr") or out["currency"]
-    m = re.search(TAX_PAT, t, re.I|re.S);
-    if m: out["tax"], out["currency"] = m.group("amt").replace(",",""), m.group("curr") or out["currency"]
-    m = re.search(TOTAL_PAT, t, re.I|re.S);
-    if m:
-        out["total"], out["currency"] = m.group("amt").replace(",", ""), m.group("curr") or out["currency"]
-    if out["currency"] in ["$", "C$", "€", "£"]:
-        out["currency"] = {"$":"USD", "C$":"CAD", "€":"EUR", "£":"GBP"}[out["currency"]]
-    return out
-
-def normalize_kv_from_donut(parsed: dict):
-    out = {k: None for k in ["invoice_number","invoice_date","po_number","subtotal","tax","total","currency"]}
-
-    def search_keys(obj, key_list):
-        if isinstance(obj, dict):
-            for k, v in obj.items():
-                kl = k.lower()
-                if any(kk in kl for kk in key_list):
-                    return v if isinstance(v, str) else None
-                found = search_keys(v, key_list)
-                if found is not None:
-                    return found
-        elif isinstance(obj, list):
-            for it in obj:
-                found = search_keys(it, key_list)
-                if found is not None:
-                    return found
-        return None
-
-    mapping = {
-        "invoice_number": ["invoice_number","invoice no","invoice_no","invoice","inv_no"],
-        "invoice_date":   ["invoice_date","date","bill_date"],
-        "po_number":      ["po_number","po","purchase_order"],
-        "subtotal":       ["subtotal","sub_total"],
-        "tax":            ["tax","gst","vat","hst"],
-        "total":          ["total","amount_total","amount_due","grand_total"],
-    }
-    for k, keys in mapping.items():
-        val = search_keys(parsed, keys)
-        if isinstance(val, str):
-            out[k] = val.strip()
-
-    txt = json.dumps(parsed, ensure_ascii=False)
-    m = re.search(r"(USD|CAD|EUR|GBP|\$|C\$|€|£)", txt, re.I)
-    if m:
-        sym = m.group(1)
-        out["currency"] = {"$":"USD","C$":"CAD","€":"EUR","£":"GBP"}.get(sym, sym.upper())
+def load_models():
+    sentence_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    json_converter = pipeline("text2text-generation", model="yahyakhoder/MD2JSON-T5-small-V1")
+    return sentence_model, json_converter
+
+sentence_model, json_converter = load_models()
+
+# ----------------------------- Helpers -----------------------------
+def extract_candidates(text: str) -> Dict[str, str]:
+    """
+    Very simple key:value candidate extraction from lines containing a colon.
+    Extend with regexes for phones/emails/totals as needed.
+    """
+    out: Dict[str, str] = {}
+    for raw in text.splitlines():
+        line = raw.strip()
+        if not line or ":" not in line:
+            continue
+        k, v = line.split(":", 1)
+        k, v = k.strip(), v.strip()
+        if k and v:
+            out[k] = v
     return out
 
-def items_from_words_simple(tsv: pd.DataFrame) -> pd.DataFrame:
-    HEAD_CANDIDATES = ["description","item","qty","quantity","price","unit","rate","amount","total"]
-    if tsv.empty:
-        return pd.DataFrame()
-
-    lines = []
-    for (b,p,l), g in tsv.groupby(["block_num","par_num","line_num"]):
-        text = " ".join([w for w in g["text"].astype(str).tolist() if w.strip()])
-        if text.strip():
-            lines.append({
-                "block_num": b, "par_num": p, "line_num": l,
-                "text": text.lower(),
-                "top": g["top"].min(), "bottom": (g["top"]+g["height"]).max(),
-                "left": g["left"].min(), "right": (g["left"]+g["width"]).max()
-            })
-    L = pd.DataFrame(lines)
-    if L.empty:
-        return pd.DataFrame()
-    L["header_score"] = L["text"].apply(lambda s: sum(1 for h in HEAD_CANDIDATES if h in s))
-    hdrs = L[L["header_score"] >= 2].sort_values(["header_score","top"], ascending=[False,True])
-    if hdrs.empty:
-        return pd.DataFrame()
-
-    H = hdrs.iloc[0]
-    header_top, header_bottom = H["top"], H["bottom"]
-
-    header_words = tsv[(tsv["top"] >= header_top - 5) & ((tsv["top"] + tsv["height"]) <= header_bottom + 5)]
-    header_words = header_words.sort_values("left")
-    if header_words.empty:
-        return pd.DataFrame()
-    xs = header_words["left"].tolist()
-    hdr_tokens = [t.lower() for t in header_words["text"].tolist()]
-
-    below = tsv[tsv["top"] > header_bottom + 5].copy()
-    totals_mask = below["text"].str.lower().str.contains(
-        r"(sub\s*total|amount\s*due|total|grand\s*total|balance)", regex=True, na=False
+def map_to_schema_headers(candidates: Dict[str, str],
+                          static_headers: List[str],
+                          thresh: float) -> Dict[str, str]:
+    if not candidates:
+        return {}
+    cand_keys = list(candidates.keys())
+    # Embed all at once
+    cand_emb = sentence_model.encode(cand_keys, normalize_embeddings=True)
+    head_emb = sentence_model.encode(static_headers, normalize_embeddings=True)
+    cos = util.cos_sim(torch.tensor(cand_emb), torch.tensor(head_emb)).cpu().numpy()  # [Nc, Nh]
+    mapped: Dict[str, str] = {}
+    for i, key in enumerate(cand_keys):
+        j = int(np.argmax(cos[i]))
+        score = float(cos[i][j])
+        if score >= thresh:
+            mapped[static_headers[j]] = candidates[key]
+    return mapped
+
+def build_prompt(schema_text: str, invoice_text: str, mapped_hints: Dict[str, str]) -> str:
+    """
+    Construct a strong instruction prompt for MD2JSON.
+    Includes your schema, strict rules, raw text, and optional hints from semantic mapping.
+    """
+    # Your strict instruction block:
+    instruction = (
+        'Use this schema:\n'
+        '{\n'
+        '  "invoice_header": {\n'
+        '    "car_number": "string or null",\n'
+        '    "shipment_number": "string or null",\n'
+        '    "shipping_point": "string or null",\n'
+        '    "currency": "string or null",\n'
+        '    "invoice_number": "string or null",\n'
+        '    "invoice_date": "string or null",\n'
+        '    "order_number": "string or null",\n'
+        '    "customer_order_number": "string or null",\n'
+        '    "our_order_number": "string or null",\n'
+        '    "sales_order_number": "string or null",\n'
+        '    "purchase_order_number": "string or null",\n'
+        '    "order_date": "string or null",\n'
+        '    "supplier_name": "string or null",\n'
+        '    "supplier_address": "string or null",\n'
+        '    "supplier_phone": "string or null",\n'
+        '    "supplier_email": "string or null",\n'
+        '    "supplier_tax_id": "string or null",\n'
+        '    "customer_name": "string or null",\n'
+        '    "customer_address": "string or null",\n'
+        '    "customer_phone": "string or null",\n'
+        '    "customer_email": "string or null",\n'
+        '    "customer_tax_id": "string or null",\n'
+        '    "ship_to_name": "string or null",\n'
+        '    "ship_to_address": "string or null",\n'
+        '    "bill_to_name": "string or null",\n'
+        '    "bill_to_address": "string or null",\n'
+        '    "remit_to_name": "string or null",\n'
+        '    "remit_to_address": "string or null",\n'
+        '    "tax_id": "string or null",\n'
+        '    "tax_registration_number": "string or null",\n'
+        '    "vat_number": "string or null",\n'
+        '    "payment_terms": "string or null",\n'
+        '    "payment_method": "string or null",\n'
+        '    "payment_reference": "string or null",\n'
+        '    "bank_account_number": "string or null",\n'
+        '    "iban": "string or null",\n'
+        '    "swift_code": "string or null",\n'
+        '    "total_before_tax": "string or null",\n'
+        '    "tax_amount": "string or null",\n'
+        '    "tax_rate": "string or null",\n'
+        '    "shipping_charges": "string or null",\n'
+        '    "discount": "string or null",\n'
+        '    "total_due": "string or null",\n'
+        '    "amount_paid": "string or null",\n'
+        '    "balance_due": "string or null",\n'
+        '    "due_date": "string or null",\n'
+        '    "invoice_status": "string or null",\n'
+        '    "reference_number": "string or null",\n'
+        '    "project_code": "string or null",\n'
+        '    "department": "string or null",\n'
+        '    "contact_person": "string or null",\n'
+        '    "notes": "string or null",\n'
+        '    "additional_info": "string or null"\n'
+        '  },\n'
+        '  "line_items": [\n'
+        '    {\n'
+        '      "quantity": "string or null",\n'
+        '      "units": "string or null",\n'
+        '      "description": "string or null",\n'
+        '      "footage": "string or null",\n'
+        '      "price": "string or null",\n'
+        '      "amount": "string or null",\n'
+        '      "notes": "string or null"\n'
+        '    }\n'
+        '  ]\n'
+        '}\n'
+        'If a field is missing for a line item or header, use null. '
+        'Do not invent fields. Do not add any header or shipment data to any line item. '
+        'Return ONLY the JSON object, no explanation.\n'
+    )
+    hint_block = ""
+    if mapped_hints:
+        # Provide soft hints to help the T5 model place values correctly (optional but useful)
+        hints_serialized = " ".join([f"#{k}: {v}" for k, v in mapped_hints.items()])
+        hint_block = f"\nHints:\n{hints_serialized}\n"
+
+    return (
+        instruction +
+        "\nInvoice Text:\n" +
+        invoice_text.strip() +
+        hint_block
     )
-    if totals_mask.any():
-        stop_y = below.loc[totals_mask, "top"].min()
-        below = below[below["top"] < stop_y - 4]
-    if below.empty:
-        return pd.DataFrame()
-
-    rows = []
-    for (b,p,l), g in below.groupby(["block_num","par_num","line_num"]):
-        g = g.sort_values("left")
-        buckets = {i:[] for i in range(len(xs))}
-        for _, w in g.iterrows():
-            if not str(w["text"]).strip():
-                continue
-            idx = int(np.abs(np.array(xs) - w["left"]).argmin())
-            buckets[idx].append(str(w["text"]))
-        vals = [" ".join(buckets[i]).strip() for i in range(len(xs))]
-        if any(vals):
-            rows.append(vals)
-    if not rows:
-        return pd.DataFrame()
-
-    df_rows = pd.DataFrame(rows).fillna("")
-    names = []
-    for i in range(df_rows.shape[1]):
-        wl = hdr_tokens[i] if i < len(hdr_tokens) else f"col_{i}"
-        if "desc" in wl or wl in ["item","description"]:
-            names.append("description")
-        elif wl in ["qty","quantity"]:
-            names.append("quantity")
-        elif "unit" in wl or "rate" in wl or "price" in wl:
-            names.append("unit_price")
-        elif "amount" in wl or "total" in wl:
-            names.append("line_total")
-        else:
-            names.append(f"col_{i}")
-    df_rows.columns = names
-    df_rows = df_rows[~(df_rows.fillna("").apply(lambda r: "".join(r.values), axis=1).str.strip()=="")]
-    return df_rows.reset_index(drop=True)
-
-# ----------------------------- App -----------------------------
-st.title("Invoice Extraction — Donut (public) + Tesseract tables")
-
-up = st.file_uploader("Upload an invoice (PDF/JPG/PNG)", type=["pdf","png","jpg","jpeg"])
-if not up:
-    st.info("Upload a scanned invoice to begin.")
-    st.stop()
-
-with st.spinner(f"Loading model '{model_id}' from Hugging Face…"):
-    processor, donut_model = load_donut(model_id)
-
-pages = load_pages(up.read(), up.name)
-page_idx = 0
-if len(pages) > 1:
-    page_idx = st.number_input("Page", 1, len(pages), 1) - 1
-img = pages[page_idx]
-
-col1, col2 = st.columns([1.1, 1.3], gap="large")
-
-with col1:
-    st.subheader("Preview")
-    st.image(img, use_column_width=True)
-    det_img = preprocess_for_detection(img)
-    with st.expander("Detection view (preprocessed for boxes)"):
-        st.image(det_img, use_column_width=True)
-
-with col2:
-    st.subheader("OCR & Extraction")
-
-    with st.spinner("Running Donut…"):
-        seq, parsed = donut_infer(img, processor, donut_model, task_prompt)
-
-    if parsed:
-        key_fields = normalize_kv_from_donut(parsed)
-        donut_payload = parsed
-    else:
-        key_fields = parse_fields_regex(seq)
-        donut_payload = {"generated_text": seq}
-
-    k1,k2,k3 = st.columns(3)
-    with k1:
-        st.write(f"**Invoice #:** {key_fields.get('invoice_number') or '—'}")
-        st.write(f"**Invoice Date:** {key_fields.get('invoice_date') or '—'}")
-    with k2:
-        st.write(f"**PO #:** {key_fields.get('po_number') or '—'}")
-        st.write(f"**Subtotal:** {key_fields.get('subtotal') or '—'}")
-    with k3:
-        st.write(f"**Tax:** {key_fields.get('tax') or '—'}")
-        tot = key_fields.get('total') or '—'
-        cur = key_fields.get('currency') or ''
-        st.write(f"**Total:** {tot} {cur}".strip())
-
-    with st.spinner("Detecting words with Tesseract (for table)…"):
-        tsv = pytesseract.image_to_data(det_img, lang=det_lang, output_type=Output.DATAFRAME)
-        tsv = tsv.dropna(subset=["text"]).reset_index(drop=True)
-        tsv["x2"] = tsv["left"] + tsv["width"]
-        tsv["y2"] = tsv["top"] + tsv["height"]
-
-    st.markdown("**Line Items**")
-    items = items_from_words_simple(tsv)
-    if items.empty:
-        st.caption("No line items confidently detected.")
-    else:
-        st.dataframe(items, use_container_width=True)
 
-    result = {
-        "file": up.name,
-        "page": page_idx + 1,
-        "key_fields": key_fields,
-        "items": items.to_dict(orient="records") if not items.empty else [],
-        "donut_raw": donut_payload,
-    }
-    st.download_button("Download JSON", data=json.dumps(result, indent=2),
-                       file_name="invoice_extraction.json", mime="application/json")
-    if not items.empty:
-        st.download_button("Download Items CSV", data=items.to_csv(index=False),
-                           file_name="invoice_items.csv", mime="text/csv")
+def strict_json_or_fallback(text: str, mapped: Dict[str, str]) -> Dict:
+    """Try to parse model output as JSON. If invalid, return a schema-shaped fallback using mapped values."""
+    # Try full parse
+    try:
+        return json.loads(text)
+    except Exception:
+        pass
+    # Try to extract the biggest {...} block
+    start = text.find("{")
+    end = text.rfind("}")
+    if start != -1 and end != -1 and end > start:
+        try:
+            return json.loads(text[start:end+1])
+        except Exception:
+            pass
+    # Fallback: shape like schema and fill mapped header fields; empty line_items
+    fallback = json.loads(json.dumps(SCHEMA_JSON))  # deep copy
+    for k, v in mapped.items():
+        if k in fallback["invoice_header"]:
+            fallback["invoice_header"][k] = v
+    # Leave line_items as a single template row with nulls
+    return fallback
+
+# ----------------------------- UI: paste text -----------------------------
+invoice_text = st.text_area(
+    "Paste the invoice text here (unstructured).",
+    height=320,
+    placeholder="Paste the invoice content (OCR/plain text) ..."
+)
 
-    if show_boxes:
-        st.caption("First 20 Tesseract word boxes")
-        st.dataframe(tsv[["left","top","width","height","text","conf"]].head(20), use_container_width=True)
+colA, colB = st.columns([1,1])
+
+with colA:
+    if st.button("Generate JSON", type="primary", use_container_width=True):
+        if not invoice_text.strip():
+            st.error("Please paste the invoice text first.")
+            st.stop()
+
+        with st.spinner("Extracting candidates..."):
+            candidates = extract_candidates(invoice_text)
+
+        with st.spinner("Semantic mapping with all-MiniLM-L6-v2..."):
+            mapped = map_to_schema_headers(candidates, STATIC_HEADERS, threshold)
+
+        if show_intermediates:
+            st.subheader("Candidates (first 20)")
+            st.json(dict(list(candidates.items())[:20]))
+            st.subheader("Mapped headers")
+            st.json(mapped)
+
+        with st.spinner("Generating structured JSON with MD2JSON-T5-small-V1..."):
+            prompt = build_prompt(json.dumps(SCHEMA_JSON, indent=2), invoice_text, mapped)
+            gen = pipeline("text2text-generation", model="yahyakhoder/MD2JSON-T5-small-V1")
+            out = gen(prompt, max_new_tokens=max_new_tokens)[0]["generated_text"]
+            final_json = strict_json_or_fallback(out, mapped)
+
+        st.subheader("Final JSON")
+        st.json(final_json)
+        st.download_button("Download JSON", data=json.dumps(final_json, indent=2),
+                           file_name="invoice.json", mime="application/json", use_container_width=True)
+
+with colB:
+    st.caption("Tips")
+    st.markdown(
+        "- You can tune the **semantic threshold** in the sidebar to make header mapping more/less strict.\n"
+        "- If the T5 model outputs invalid JSON, the app returns a **schema-shaped fallback** filled from the mapped fields.\n"
+        "- Extend `extract_candidates()` to add regexes for emails, phones, totals, etc., to boost recall."
+    )