app.py

"""nutrient-document-decision demo.

One tab per task, each with a UI suited to that task and a live-reactive panel showing the exact
`state`/`questions` JSON that request would send -- it updates as you type/upload, not just after
running. A separate Custom tab is the free-form request builder for anything the three task tabs
don't cover; images there are referenced by the server-side temp path they were uploaded to (shown
for you to copy into the JSON wherever an image belongs), so any number of images, anywhere in the
structure, just works.
"""
import base64
import json
import tempfile
from pathlib import Path

import fitz  # pymupdf
import gradio as gr
import pytesseract
import spaces
from PIL import Image

import serve

EXAMPLES_DIR = Path(__file__).resolve().parent / "examples"
_loaded = False

BRAND_CSS = """
:root { --brand-accent: #1d5fa8; }
.gradio-container { font-family: -apple-system, "Helvetica Neue", Arial, sans-serif; }
button.primary, .primary { background: var(--brand-accent) !important; border-color: var(--brand-accent) !important; }
a { color: var(--brand-accent); }
.json-box textarea, .json-box .cm-editor { font-size: 12.5px; }
.json-box .cm-editor { max-height: 400px; overflow-y: auto; }
.json-box textarea { max-height: 400px; overflow-y: auto !important; }
"""


def _ensure_loaded():
    global _loaded
    if not _loaded:
        serve.load(device="cuda")
        _loaded = True


# --------------------------------------------------------------------------- file/PDF handling

def _pdf_pages(path, max_pages=40):
    doc = fitz.open(path)
    paths = []
    for i, page in enumerate(doc):
        if i >= max_pages:
            break
        pix = page.get_pixmap(matrix=fitz.Matrix(150 / 72, 150 / 72))
        img = Image.frombytes("RGB", (pix.width, pix.height), pix.samples)
        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        img.save(tmp.name)
        paths.append(tmp.name)
    doc.close()
    return paths


def _tiff_pages(path, max_pages=40):
    im = Image.open(path)
    n = getattr(im, "n_frames", 1)
    if n <= 1:
        return [path]
    paths = []
    for i in range(min(n, max_pages)):
        im.seek(i)
        tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
        im.convert("RGB").save(tmp.name)
        paths.append(tmp.name)
    return paths


def _expand_to_pages(files):
    """files: list of uploaded file paths (images and/or PDFs). Returns an ordered list of
    single-page image paths, expanding any PDF or multi-page TIFF into one entry per page."""
    out = []
    for f in files or []:
        if not f:
            continue
        low = str(f).lower()
        if low.endswith(".pdf"):
            out.extend(_pdf_pages(f))
        elif low.endswith((".tif", ".tiff")):
            out.extend(_tiff_pages(f))
        else:
            out.append(f)
    return out


def _as_list(files):
    if files is None:
        return []
    return files if isinstance(files, list) else [files]


# --------------------------------------------------------------------------- Grounding tab

def build_grounding_request(source, claim):
    state = {"source": source or "", "claim": claim or ""}
    questions = {
        "grounded": {
            "type": "noul",
            "instructions": "The claim is fully and explicitly supported by the source, with no "
            "contradiction and nothing invented.",
        }
    }
    return state, questions


def grounding_preview(source, claim):
    state, questions = build_grounding_request(source, claim)
    return json.dumps({"state": state, "questions": questions}, indent=2)


@spaces.GPU(duration=60)
def _grounding_run(source, claim):
    state, questions = build_grounding_request(source, claim)
    _ensure_loaded()
    return serve.answer(state, questions)


def _grounding_render(result):
    p = result["grounded"]["noul"]
    verdict = "✅ grounded" if p >= 0.5 else "❌ not grounded"
    bar = "█" * round(p * 20) + "░" * (20 - round(p * 20))
    return f"### {verdict}\n\n`{bar}` **P(grounded) = {p:.3f}**"


def run_grounding(source, claim):
    if not (source or "").strip() or not (claim or "").strip():
        raise gr.Error("Enter both a source and a claim.")
    yield "⏳ Waking up… cold-starting a shared GPU (~30–60s)…", "", ""
    result = _grounding_run(source, claim)
    yield "", _grounding_render(result), json.dumps(result, indent=2)


# --------------------------------------------------------------------------- Document Classification tab

def _parse_labels(text):
    """`label | description` per line (description optional) -- same convention as
    document-classification-demo. Returns an ordered {label: description-or-label} dict."""
    out = {}
    for line in (text or "").splitlines():
        label, _, desc = line.partition("|")
        label, desc = label.strip(), desc.strip()
        if label:
            out[label] = desc or label
    return out


DC_DEFAULT_LABELS = (
    "invoice | an itemized bill listing goods or services and a total due\n"
    "tax form | a government tax filing with numbered lines and boxes\n"
    "receipt | a short proof-of-purchase slip with a total\n"
    "letter | correspondence with a salutation, body, and signature"
)

_ocr_cache: dict[str, str] = {}


def _ocr_text(path):
    """Tesseract OCR, cached per page path -- this is the OCR the published checkpoint expects
    alongside the image (it was trained with OCR present)."""
    if path not in _ocr_cache:
        try:
            _ocr_cache[path] = pytesseract.image_to_string(Image.open(path)).strip()
        except Exception:
            _ocr_cache[path] = ""
    return _ocr_cache[path]


def build_dc_request(files, labels_text):
    pages = _expand_to_pages(_as_list(files))
    criteria = _parse_labels(labels_text)
    state = {
        f"page_{i + 1}": {"image": {"type": "image", "source": p}, "text": _ocr_text(p)}
        for i, p in enumerate(pages)
    }
    questions = {
        f"page_{i + 1}_doctype": {
            "type": "choice",
            "instructions": f"Identify the type of page_{i + 1}.",
            "criteria": criteria,
        }
        for i in range(len(pages))
    }
    return state, questions, pages


def dc_preview(files, labels_text):
    state, questions, _ = build_dc_request(files, labels_text)
    return json.dumps({"state": state, "questions": questions}, indent=2)


@spaces.GPU(duration=90)
def _dc_run(files, labels_text):
    state, questions, pages = build_dc_request(files, labels_text)
    if not pages:
        raise gr.Error("Upload at least one page (image or PDF).")
    if len(_parse_labels(labels_text)) < 2:
        raise gr.Error("Enter at least two candidate classes (one per line: label | description).")
    _ensure_loaded()
    return serve.answer(state, questions), len(pages)


def _dc_render(result, n_pages):
    blocks = []
    for i in range(n_pages):
        r = result[f"page_{i + 1}_doctype"]
        ranked = sorted(r["probabilities"].items(), key=lambda kv: -kv[1])
        rows = []
        for rank, (label, p) in enumerate(ranked, 1):
            pct = max(0.0, min(1.0, p)) * 100
            star = " ✅" if rank == 1 else ""
            rows.append(
                f"<tr><td style='text-align:right;padding:4px 10px;color:#5c6773'>{rank}</td>"
                f"<td style='padding:4px 10px'><b>{label}</b>{star}</td>"
                f"<td style='padding:4px 10px;width:160px'>"
                f"<div style='background:#e6ecf5;border-radius:4px;height:14px'>"
                f"<div style='background:#1d5fa8;height:14px;border-radius:4px;width:{pct:.0f}%'></div></div></td>"
                f"<td style='text-align:right;padding:4px 10px;font-variant-numeric:tabular-nums'>{p:.3f}</td></tr>"
            )
        blocks.append(
            f"<div style='font:600 13px system-ui;margin:12px 0 4px'>Page {i + 1}</div>"
            "<table style='border-collapse:collapse;font:14px system-ui'><tbody>" + "".join(rows) + "</tbody></table>"
        )
    return "".join(blocks)


def run_dc(files, labels_text):
    yield "⏳ Waking up… OCR + cold-starting a shared GPU (~30–60s)…", "", ""
    result, n_pages = _dc_run(files, labels_text)
    yield "", _dc_render(result, n_pages), json.dumps(result, indent=2)


# --------------------------------------------------------------------------- Document Split tab

_thumb_cache: dict[str, str] = {}


def _thumb_uri(path, height=150):
    import io

    if path not in _thumb_cache:
        im = Image.open(path).convert("RGB")
        im.thumbnail((height * 2, height * 2))
        buf = io.BytesIO()
        im.save(buf, "PNG")
        _thumb_cache[path] = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
    return _thumb_cache[path]


def _thumbstrip(pages):
    parts = ["<div style='display:flex;gap:8px;overflow-x:auto;padding:8px 0'>"]
    for i, p in enumerate(pages):
        parts.append(
            f"<div style='text-align:center;flex:0 0 auto'><img src='{_thumb_uri(p)}' "
            f"style='height:160px;border:1px solid #ccc;border-radius:4px'/>"
            f"<div style='font:11px system-ui;color:#888'>page {i + 1}</div></div>"
        )
    parts.append("</div>")
    return "".join(parts)


def build_ds_request(files):
    pages = _expand_to_pages(_as_list(files))
    state = {f"page_{i + 1}": {"image": {"type": "image", "source": p}} for i, p in enumerate(pages)}
    questions = {
        f"page_{i + 1}_boundary": {
            "type": "noul",
            "instructions": f"page_{i + 1} begins a new, separate document from page_{i} -- "
            "they are not pages of the same document.",
        }
        for i in range(1, len(pages))
    }
    return state, questions, pages


def ds_on_upload(files):
    _, _, pages = build_ds_request(files)
    return _thumbstrip(pages) if pages else ""


def ds_preview(files):
    state, questions, _ = build_ds_request(files)
    return json.dumps({"state": state, "questions": questions}, indent=2)


@spaces.GPU(duration=90)
def _ds_run(files):
    state, questions, pages = build_ds_request(files)
    if len(pages) < 2:
        raise gr.Error("Upload at least two pages (images and/or a multi-page PDF).")
    _ensure_loaded()
    return serve.answer(state, questions), pages


def _ds_render(documents, pages):
    """Mirrors doc-split-demo's card layout: one bordered card per document, an uppercase
    accent-colored header naming the page range and split confidence, and each page's own
    thumbnail inside the card."""
    html = [
        f"<div style='font:600 15px system-ui;margin:0 0 12px'>Split into <b>{len(documents)}</b> "
        f"document(s) across {len(pages)} page(s)</div>"
    ]
    for k, d in enumerate(documents, 1):
        conf = d["boundary_confidence"]
        if conf is None:
            dot, label = "#5c6773", "forced start (p1)"
        else:
            dot = "#0f7a58" if conf >= 0.8 else ("#9a6a12" if conf >= 0.5 else "#a83a3a")
            label = f"split confidence {conf:.0%}"
        group = d["pages"]
        html.append(
            "<div style='border:1px solid #d5deea;border-radius:10px;padding:12px;margin:0 0 12px'>"
            "<div style='font:600 12px system-ui;letter-spacing:.05em;color:#1d5fa8;"
            "text-transform:uppercase;margin-bottom:8px'>"
            f"Document {k} · pages {group[0]}–{group[-1]} ({len(group)}p) · "
            f"<span style='color:{dot}'>&#9679; {label}</span></div>"
            "<div style='display:flex;gap:8px;overflow-x:auto'>"
        )
        for i in group:
            html.append(
                f"<div style='text-align:center;flex:0 0 auto'>"
                f"<img src='{_thumb_uri(pages[i - 1])}' style='height:150px;border:1px solid #ccc;border-radius:4px'/>"
                f"<div style='font:11px system-ui;color:#5c6773'>p{i}</div></div>"
            )
        html.append("</div></div>")
    return "".join(html)


def run_ds(files):
    yield "⏳ Waking up… cold-starting a shared GPU (~30–60s)…", "", ""
    result, pages = _ds_run(files)
    boundary = [0.0] + [result[f"page_{i + 1}_boundary"]["noul"] for i in range(1, len(pages))]
    docs, cur = [], []
    for i, p in enumerate(boundary):
        if p >= 0.5 and cur:
            docs.append(cur)
            cur = []
        cur.append(i + 1)
    if cur:
        docs.append(cur)
    documents = [{"pages": group,
                  "boundary_confidence": None if group[0] == 1 else round(boundary[group[0] - 1], 6)}
                 for group in docs]
    out = {"documents": documents, "model_response": result}
    yield "", _ds_render(documents, pages), json.dumps(out, indent=2)


# --------------------------------------------------------------------------- Custom tab

def files_to_paths(files):
    """Gradio already saved each upload to a server-side temp path -- just hand that path back for
    the user to paste as a `"source"` value. A multi-page TIFF expands into one path per page
    (each saved as its own temp PNG), same as the task tabs. No encoding needed, and it works for
    any number of images anywhere in the JSON structure. Rendered as side-by-side cards (click a
    path to select it for copying) rather than one long stacked list."""
    if not files:
        return ""
    cards = []
    for f in _as_list(files):
        name = Path(f).name
        pages = _expand_to_pages([f])
        for i, p in enumerate(pages):
            label = name if len(pages) == 1 else f"{name} (page {i + 1})"
            cards.append(
                "<div style='flex:0 0 auto;min-width:220px'>"
                f"<div style='font:600 12px system-ui;margin-bottom:4px'>{label}</div>"
                f"<input readonly value='{p}' onclick='this.select()' "
                "style='width:100%;font:12px ui-monospace,monospace;padding:4px 6px;"
                "border:1px solid #ccc;border-radius:4px'/></div>"
            )
    return "<div style='display:flex;gap:14px;overflow-x:auto;padding:4px 0'>" + "".join(cards) + "</div>"


@spaces.GPU(duration=60)
def _custom_run(state_text, questions_text):
    try:
        state = json.loads(state_text)
        questions = json.loads(questions_text)
    except json.JSONDecodeError as e:
        raise gr.Error(f"Invalid JSON: {e}")
    _ensure_loaded()
    return serve.answer(state, questions)


def run_custom(state_text, questions_text):
    yield "⏳ Waking up… first run downloads the model and cold-starts a shared GPU (~30–60s)…"
    yield json.dumps(_custom_run(state_text, questions_text), indent=2)


CUSTOM_DEFAULT_STATE = json.dumps(
    {"source": "The invoice total is $482.10, due June 1, 2026.", "claim": "The invoice is due June 1, 2026."},
    indent=2,
)
CUSTOM_DEFAULT_QUESTIONS = json.dumps(
    {"grounded": {"type": "noul", "instructions": "The claim is fully supported by the source."}}, indent=2
)

LINKS = (
    "→ [model card](https://huggingface.co/nutrientdocs/nutrient-document-decision) · "
    "[grounding leaderboard](https://huggingface.co/spaces/nutrientdocs/grounding-leaderboard) · "
    "[document classification leaderboard](https://huggingface.co/spaces/nutrientdocs/document-classification-leaderboard) · "
    "[open-vocabulary leaderboard](https://huggingface.co/spaces/nutrientdocs/doc-openvocab-leaderboard) · "
    "[document split leaderboard](https://huggingface.co/spaces/nutrientdocs/doc-split-leaderboard)"
)
GET_ACCESS = (
    "> ### 📩 Get access\n"
    "> `nutrient-document-decision` is commercial and its weights are not downloadable here. "
    "To run it on-prem — **contact Nutrient: [nutrient.io/contact-sales](https://www.nutrient.io/contact-sales/).**"
)
ABOUT_AUTHOR = (
    '## About the author\n'
    '<a href="https://nutrient.io/">'
    '<img src="https://avatars2.githubusercontent.com/u/1527679?v=3&s=200" height="80" /></a>\n\n'
    "This demo is maintained and funded by [Nutrient](https://nutrient.io/) — "
    "The deterministic document infrastructure enterprises run their highest-stakes workflows on: "
    "replayable output, clear exceptions, and full audit trails on the messy, regulated documents where AI alone breaks."
)

with gr.Blocks(title="Document Decision Demo") as demo:
    gr.Markdown(
        "# nutrient-document-decision\n"
        "A decision model for documents -- in the spirit of general decision-making systems like "
        "Jev, but focused on document and multimodal understanding. Ask it typed questions about a "
        "document in a single pass; grounding, document classification, and document-split are the "
        "tasks shown below, not the limit of what it can answer -- it has generalized zero-shot to "
        "related judgment tasks too, like comparing OCR output quality. Every tab's right-hand panel "
        "is the exact `state`/`questions` request it would send, live as you edit.\n\n" + LINKS
    )

    with gr.Tab("Grounding"):
        with gr.Row():
            with gr.Column():
                g_source = gr.Textbox(label="Source", lines=6, value="The invoice total is $482.10, due June 1, 2026.")
                g_claim = gr.Textbox(label="Claim", lines=3, value="The invoice is due June 1, 2026.")
            with gr.Column():
                g_req = gr.Code(label="Request JSON", language="json", lines=12, elem_classes=["json-box"])
        g_btn = gr.Button("Check grounding", variant="primary")
        g_status = gr.Markdown()
        with gr.Row():
            g_render = gr.Markdown(label="Result")
            g_out = gr.Code(label="Response JSON", language="json", lines=10, elem_classes=["json-box"])
        for comp in (g_source, g_claim):
            comp.change(grounding_preview, [g_source, g_claim], g_req)
        demo.load(grounding_preview, [g_source, g_claim], g_req)
        g_btn.click(run_grounding, [g_source, g_claim], [g_status, g_render, g_out])

    with gr.Tab("Document Classification"):
        with gr.Row():
            with gr.Column():
                dc_files = gr.File(label="Page(s) — images, TIFF, and/or a PDF",
                                    file_types=["image", ".pdf", ".tif", ".tiff"], file_count="multiple")
                dc_labels = gr.Textbox(label="Candidate classes (one per line: label | description)",
                                        value=DC_DEFAULT_LABELS, lines=5)
                gr.Markdown("<sub>OCR text is extracted automatically (Tesseract) — the published "
                            "checkpoint was trained with OCR present.</sub>")
            with gr.Column():
                dc_req = gr.Code(label="Request JSON", language="json", lines=12, elem_classes=["json-box"])
        dc_btn = gr.Button("Classify", variant="primary")
        dc_status = gr.Markdown()
        with gr.Row():
            dc_render = gr.HTML(label="Result")
            dc_out = gr.Code(label="Response JSON", language="json", lines=10, elem_classes=["json-box"])
        for comp in (dc_files, dc_labels):
            comp.change(dc_preview, [dc_files, dc_labels], dc_req)
        demo.load(dc_preview, [dc_files, dc_labels], dc_req)
        dc_btn.click(run_dc, [dc_files, dc_labels], [dc_status, dc_render, dc_out], show_progress="full")

    with gr.Tab("Document Split"):
        with gr.Row():
            with gr.Column():
                ds_files = gr.File(label="Page images, TIFF, and/or PDF(s), in order",
                                    file_types=["image", ".pdf", ".tif", ".tiff"], file_count="multiple")
                ds_preview_html = gr.HTML()
            with gr.Column():
                ds_req = gr.Code(label="Request JSON", language="json", lines=12, elem_classes=["json-box"])
        ds_btn = gr.Button("Split", variant="primary")
        ds_status = gr.Markdown()
        with gr.Row():
            ds_render = gr.HTML(label="Result")
            ds_out = gr.Code(label="Response JSON", language="json", lines=10, elem_classes=["json-box"])
        ds_files.change(ds_on_upload, [ds_files], ds_preview_html)
        ds_files.change(ds_preview, [ds_files], ds_req)
        ds_btn.click(run_ds, [ds_files], [ds_status, ds_render, ds_out], show_progress="full")

    with gr.Tab("Custom"):
        gr.Markdown(
            "Free-form request builder — for question types or state shapes the tabs above don't cover. "
            "Upload image(s) below to get their server-side temp path, then paste it as any `\"source\"` "
            "value in the state JSON (as many images, anywhere in the structure, as you like)."
        )
        c_files = gr.File(label="Image(s)/TIFF → temp path", file_types=["image", ".tif", ".tiff"],
                          file_count="multiple")
        gr.Markdown("<sub>Temp paths — click one to select it for copying into the state JSON below.</sub>")
        c_paths = gr.HTML()
        with gr.Row():
            c_state = gr.Code(label="state (JSON)", language="json", lines=12, value=CUSTOM_DEFAULT_STATE, elem_classes=["json-box"])
            c_questions = gr.Code(label="questions (JSON)", language="json", lines=12, value=CUSTOM_DEFAULT_QUESTIONS, elem_classes=["json-box"])
        c_btn = gr.Button("Run", variant="primary")
        c_out = gr.Code(label="Response", language="json", lines=8, elem_classes=["json-box"])
        c_files.change(files_to_paths, [c_files], c_paths)
        c_btn.click(run_custom, [c_state, c_questions], c_out, show_progress="full")

    gr.Markdown(GET_ACCESS)
    gr.Markdown(ABOUT_AUTHOR)

if __name__ == "__main__":
    demo.queue().launch(theme=gr.themes.Soft(primary_hue="blue"), css=BRAND_CSS)