render.py

"""Result-panel rendering for the Polygen Qwen2-VL Space.

Pure presentation + scoring logic, kept out of `app.py` so it is unit-testable
without a Gradio runtime or a live endpoint. Everything here is honest-by-
construction:

- Token counts prefer the worker's MEASURED `n_polygen_visual_tokens` and only
  fall back to the length law when the worker did not surface it.
- Correctness uses word-boundary token matching, NOT substring matching, so a
  ground-truth token like ``"red"`` does not spuriously match ``"covered"``.
- The verdict is correctness-agreement when ground truth exists, so two
  differently-worded-but-both-correct answers read as "Same answer" instead of
  the near-always "Different" that strict string equality produces.
- Compute claims are FLOP-count reductions (attention is $O(L^2)$ in sequence
  length), never wall-clock. No "faster" framing anywhere.
"""

import re
from html import escape
from typing import Optional

# ---------------------------------------------------------------------------
# Length law + FLOP math.
# ---------------------------------------------------------------------------

# Per-tier storage compression from the SDK tier registry
# (polygen/integrations/qwen2vl/tiers.py). Fallback only: the worker echoes the
# real value as `compression_ratio`. Lower tier keeps fewer lossless tiles, so
# it compresses storage harder (T=2 = 7.5x > T=3 = 5.1x).
_TIER_STORAGE = {0: 283.0, 1: 14.3, 2: 7.5, 3: 5.1, 4: 3.9}


def bypass_n_polygen_visual_tokens(
    n_visual: int,
    tier: int,
    M: int = 4,
    K: int = 3,
) -> int:
    """Post-bypass visual-token count -- FALLBACK ONLY.

    Prefer the worker's measured `n_polygen_visual_tokens`; this reproduces the
    SDK length law for the case the worker did not surface it. Bypass keeps
    ``tier`` of ``M`` tiles at full length and collapses the remaining
    ``M - tier`` tiles to ``K`` coefficient tokens each, so
    ``L_v_eff = sum(kept tile rows) + (M - tier) * K``. Tiles come from
    ``np.array_split``, so the first ``n_visual % M`` tiles carry one extra row;
    this matters by at most 1 token versus the measured value. Verified against
    the decomp_7B points (121->94, 676->510, 900->678, 1089->820, 1369->1030).
    """
    if tier >= M:
        return n_visual
    if tier <= 0:
        return M * K
    base, rem = divmod(n_visual, M)
    sizes = [base + 1] * rem + [base] * (M - rem)
    kept = sum(sorted(sizes, reverse=True)[:tier])
    return kept + (M - tier) * K


def attention_flops_reduction_pct(n_visual: int, n_poly: int) -> int:
    """Percent reduction in visual self-attention FLOPs ($O(L^2)$).

    FLOP-count claim, not wall-clock. Returns 0 when the sequence did not
    shorten (hook path, T=4, or a missing measurement).
    """
    if n_visual <= 0 or n_poly <= 0 or n_poly >= n_visual:
        return 0
    return round(100 * (1 - (n_poly / n_visual) ** 2))


# ---------------------------------------------------------------------------
# Matching + correctness.
# ---------------------------------------------------------------------------

_WORD = re.compile(r"[a-z0-9]+")


def _tokens(s: str) -> list[str]:
    return _WORD.findall((s or "").lower())


def _norm_phrase(s: str) -> str:
    return " ".join(_tokens(s))


def is_correct(answer: str, ground_truth: Optional[list[str]]) -> Optional[bool]:
    """Score an answer against ground truth. ``None`` when no ground truth.

    Word-boundary matching, NOT substring: a single-token ground-truth variant
    must appear as a whole token (so ``"red"`` does not match ``"covered"``); a
    multi-token variant must appear as a contiguous token phrase. Keep
    ground-truth variants short and unambiguous; constrain the gallery questions
    so the answer is a single noun / number / color. Residual risk: negation
    ("there is no cat") still tokenizes ``cat`` -- the smoke test's adversarial
    wrong-answer check is the guard for that.
    """
    if not ground_truth:
        return None
    ans_tokens = set(_tokens(answer))
    ans_phrase = " ".join(_tokens(answer))
    for variant in ground_truth:
        gt = _tokens(variant)
        if not gt:
            continue
        if len(gt) == 1:
            if gt[0] in ans_tokens:
                return True
        elif " ".join(gt) in ans_phrase:
            return True
    return False


def answers_match(
    baseline: str,
    polygen: str,
    b_correct: Optional[bool] = None,
    p_correct: Optional[bool] = None,
) -> bool:
    """Did polygen preserve the baseline's answer?

    With ground truth, this is correctness-agreement: both correct reads as
    preserved even when the wording differs ("A cat" vs "It's a cat"); a
    correct/incorrect split reads as diverged. Without ground truth (user
    uploads) or when both are wrong, fall back to normalized equality then
    bidirectional containment. Containment can overmatch on long answers
    ("a cat in the garden" contains "a cat"); for VQA-short answers that is
    fine, and the smoke test confirms longer upload-style answers don't trip a
    spurious match. If they do, tighten to token-set Jaccard or first-noun
    match.
    """
    if b_correct is not None and p_correct is not None:
        if b_correct and p_correct:
            return True
        if b_correct != p_correct:
            return False
        # both wrong: did they at least produce the same wrong answer?
    a, b = _norm_phrase(baseline), _norm_phrase(polygen)
    if not a or not b:
        return False
    if a == b:
        return True
    return a in b or b in a


# ---------------------------------------------------------------------------
# HTML render.
# ---------------------------------------------------------------------------


def render_result_panel(
    *,
    baseline_answer: str,
    polygen_answer: str,
    n_visual_tokens: int,
    n_polygen_tokens: int,  # measured (worker); <= 0 means fall back to law
    splice_mode_resolved: str,  # "hook" | "bypass"
    tier: int,
    storage_compression_x: Optional[float],
    ground_truth: Optional[list[str]],  # None for uploads
) -> str:
    """Build the combined result panel HTML.

    Per-result correctness (the verdict and the per-card chips) is scored
    against ``ground_truth`` when present; uploads carry none and are unscored.
    """
    b_correct = is_correct(baseline_answer, ground_truth)
    p_correct = is_correct(polygen_answer, ground_truth)
    matched = answers_match(baseline_answer, polygen_answer, b_correct, p_correct)

    storage_x = (
        float(storage_compression_x)
        if storage_compression_x
        else _TIER_STORAGE.get(tier, 0.0)
    )

    return f"""
<div class="pg-result">
  {_verdict_bar(matched, splice_mode_resolved)}
  <div class="pg-cards">
    {_baseline_card(baseline_answer, n_visual_tokens, b_correct)}
    {
        _polygen_card(
            answer=polygen_answer,
            n_visual=n_visual_tokens,
            n_poly_measured=n_polygen_tokens,
            mode=splice_mode_resolved,
            tier=tier,
            storage_x=storage_x,
            correct=p_correct,
        )
    }
  </div>
  {
        _explainer_footer(
            mode=splice_mode_resolved,
            n_visual=n_visual_tokens,
            n_poly_measured=n_polygen_tokens,
            tier=tier,
            storage_x=storage_x,
        )
    }
</div>
""".strip()


def _resolve_n_poly(n_visual: int, n_poly_measured: int, tier: int) -> int:
    """Prefer the worker's measured value; fall back to the length law."""
    if n_poly_measured and n_poly_measured > 0:
        return int(n_poly_measured)
    return bypass_n_polygen_visual_tokens(n_visual, tier)


def _verdict_bar(matched: bool, mode: str) -> str:
    if matched:
        sub = {
            "hook": "Polygen preserved the answer with a length-preserving "
            "compressed representation.",
            "bypass": "Polygen preserved the answer with a shorter visual sequence.",
        }.get(mode, "Polygen preserved the answer with a compressed representation.")
        return f"""<div class="pg-verdict pg-verdict-match" role="status">
  <span class="pg-verdict-icon">&#10003;</span>
  <span class="pg-verdict-text"><strong>Same answer</strong></span>
  <span class="pg-verdict-sub">{escape(sub)}</span>
</div>"""
    return """<div class="pg-verdict pg-verdict-diverge" role="status">
  <span class="pg-verdict-icon">&#9888;</span>
  <span class="pg-verdict-text"><strong>Different answer</strong></span>
  <span class="pg-verdict-sub">At this compression tier the two paths
  diverged on this input. Surfaced honestly; the per-card chips show which
  one matched the ground truth.</span>
</div>"""


def _correctness_chip(correct: Optional[bool]) -> str:
    if correct is None:
        return ""
    if correct:
        return '<span class="pg-chip pg-chip-correct">&#10003; correct</span>'
    return '<span class="pg-chip pg-chip-incorrect">&#10007; incorrect</span>'


def _baseline_card(answer: str, n_visual: int, correct: Optional[bool]) -> str:
    answer_html = (
        escape(answer)
        if answer
        else '<span class="pg-answer-empty">(run a comparison)</span>'
    )
    return f"""<article class="pg-card pg-card-base">
  <header class="pg-card-header">Baseline &middot; Qwen2-VL</header>
  <div class="pg-answer">{answer_html}</div>
  <footer class="pg-meta">
    <span class="pg-chip pg-chip-neutral">{n_visual} visual tokens</span>
    {_correctness_chip(correct)}
  </footer>
</article>"""


def _polygen_card(
    *,
    answer: str,
    n_visual: int,
    n_poly_measured: int,
    mode: str,
    tier: int,
    storage_x: float,
    correct: Optional[bool],
) -> str:
    chips: list[str] = []
    if mode == "hook":
        chips.append(
            '<span class="pg-chip pg-chip-mode-hook">'
            "Hook &middot; length-preserving</span>",
        )
        chips.append(
            f'<span class="pg-chip pg-chip-tokens">'
            f"{n_visual} visual tokens preserved</span>",
        )
    elif mode == "bypass":
        n_poly = _resolve_n_poly(n_visual, n_poly_measured, tier)
        flops = attention_flops_reduction_pct(n_visual, n_poly)
        fewer = round(100 * (1 - n_poly / max(n_visual, 1)))
        chips.append(
            '<span class="pg-chip pg-chip-mode-bypass">'
            "Bypass &middot; sequence-shortening</span>",
        )
        chips.append(
            f'<span class="pg-chip pg-chip-tokens">'
            f"{n_visual} &rarr; {n_poly} visual tokens "
            f"<small>({fewer}% fewer)</small></span>",
        )
        if flops > 0:
            chips.append(
                f'<span class="pg-chip pg-chip-flops">'
                f"{flops}% fewer attention FLOPs</span>",
            )
    else:
        chips.append(
            f'<span class="pg-chip pg-chip-neutral">{n_visual} visual tokens</span>',
        )
    if storage_x and storage_x > 1.0:
        chips.append(
            f'<span class="pg-chip pg-chip-storage">'
            f"{storage_x:.1f}&times; storage compression</span>",
        )
    chips.append(_correctness_chip(correct))

    answer_html = (
        escape(answer)
        if answer
        else '<span class="pg-answer-empty">(run a comparison)</span>'
    )
    return f"""<article class="pg-card pg-card-poly">
  <header class="pg-card-header">Polygen + LoRA &middot; T={tier}</header>
  <div class="pg-answer">{answer_html}</div>
  <footer class="pg-meta">
    {"".join(c for c in chips if c)}
  </footer>
</article>"""


def _explainer_footer(
    *,
    mode: str,
    n_visual: int,
    n_poly_measured: int,
    tier: int,
    storage_x: float,
) -> str:
    if mode == "hook":
        sentence = (
            f"Polygen encoded this image's {n_visual} visual tokens as "
            f"polynomial coefficients and reconstructed them in place for the "
            f"model. Full sequence length preserved (2-D position intact); the "
            f"stored representation is about {storage_x:.1f}x smaller."
        )
    elif mode == "bypass":
        n_poly = _resolve_n_poly(n_visual, n_poly_measured, tier)
        flops = attention_flops_reduction_pct(n_visual, n_poly)
        sentence = (
            f"Polygen replaced {n_visual} visual tokens with {n_poly} "
            f"representative tokens. Same answer over a shorter sequence: about "
            f"{flops}% fewer attention FLOPs (FLOP-count, not wall-clock), and "
            f"the stored representation about {storage_x:.1f}x smaller."
        )
    else:
        sentence = (
            f"Polygen compressed the visual representation about {storage_x:.1f}x."
        )
    return f'<div class="pg-footer">{escape(sentence)}</div>'


# ---------------------------------------------------------------------------
# CSS. Prefixed `pg-` to avoid colliding with Gradio's classes. Tuned for the
# Datasent brand theme and legible in both light and dark color schemes.
# ---------------------------------------------------------------------------

CARD_CSS = """
@media (max-width:680px) { .pg-bench-row { flex-wrap:wrap; } .pg-bench-anno { flex:1 1 100% !important; white-space:normal !important; padding-left:54px; margin-top:2px; } .pg-hero-h1 { font-size:1.6em !important; } }

.pg-result { display:flex; flex-direction:column; gap:12px; }
.pg-verdict { display:flex; align-items:center; gap:10px; flex-wrap:wrap;
  padding:10px 16px; border-radius:10px; font-size:14px; line-height:1.4; }
.pg-verdict-icon { font-size:1.15em; font-weight:700; }
.pg-verdict-sub { opacity:0.8; font-size:0.92em; }
.pg-verdict-match { background:#dcfce7; color:#166534; }
.pg-verdict-diverge { background:#fef3c7; color:#92400e; }
.dark .pg-verdict-match { background:rgba(34,197,94,0.16); color:#bbf7d0; }
.dark .pg-verdict-diverge { background:rgba(245,158,11,0.16); color:#fde68a; }
.pg-cards { display:grid; grid-template-columns:1fr 1fr; gap:14px; }
@media (max-width:680px) { .pg-cards { grid-template-columns:1fr; } }
.pg-card { border:1px solid rgba(148,163,184,0.28); border-radius:14px;
  padding:20px 22px; box-shadow:0 1px 3px rgba(0,0,0,0.12);
  background:var(--block-background-fill, rgba(255,255,255,0.04)); }
.pg-card-base { border-top:3px solid #9aa0ad; }
.pg-card-poly { border-top:3px solid #3c3475; }
.pg-card-header { font-size:0.76em; font-weight:700; letter-spacing:0.12em;
  text-transform:uppercase; opacity:0.7; margin-bottom:12px; }
.pg-answer { font-size:1.15em; line-height:1.5; min-height:2.6em;
  color:#1e293b; }
.dark .pg-answer { color:#f8fafc; }
.pg-answer-empty { opacity:0.45; }
.pg-meta { display:flex; gap:8px; margin-top:16px; flex-wrap:wrap;
  align-items:center; }
.pg-chip { display:inline-flex; align-items:center; font-size:0.78em;
  font-weight:600; padding:4px 10px; border-radius:8px;
  background:rgba(148,163,184,0.16); color:inherit; line-height:1.3; }
.pg-chip small { font-weight:400; opacity:0.75; margin-left:4px; }
.pg-chip-mode-hook { background:rgba(60,52,117,0.12); color:#3c3475; }
.pg-chip-mode-bypass { background:rgba(60,52,117,0.18); color:#3c3475; }
.dark .pg-chip-mode-hook, .dark .pg-chip-mode-bypass { color:#c1cef7; }
.pg-chip-flops { background:rgba(60,52,117,0.16); color:#3c3475; font-weight:700; }
.dark .pg-chip-flops { color:#c1cef7; }
.pg-chip-storage { background:rgba(76,170,255,0.18); color:#15618f; }
.dark .pg-chip-storage { color:#9fd2ff; }
.pg-chip-correct { background:#dcfce7; color:#166534; }
.pg-chip-incorrect { background:#fef3c7; color:#92400e; }
.dark .pg-chip-correct { background:rgba(34,197,94,0.18); color:#bbf7d0; }
.dark .pg-chip-incorrect { background:rgba(245,158,11,0.18); color:#fde68a; }
.pg-footer { font-size:0.86em; opacity:0.72; line-height:1.55;
  padding:2px 2px 0; }
"""