qwen_test_runner/vision/fuse_prompt.py

"""
fuse_prompt.py — deterministic natural-language rendering of a FusedScene.

A pure function of the fused JSON: fixed clause order (counts → primary entity →
other entities by saliency → relations → shared basin → scene), fixed attribute
ordering (consensus desc → ownership confidence desc → stratum precedence), zero
randomness — `fused_prompt(scene) == fused_prompt(scene)` byte-for-byte is a unit
test. Uncertainty is RENDERED, never guessed away: basin items become "One of
{candidates} has {attribute}." LLM smoothing is deliberately not here — if ever
wanted it is a separate additional dataset column, so this one stays trustworthy.
"""

from __future__ import annotations

from .strata import STRATUM_PRECEDENCE

_NUM_WORDS = {1: "one", 2: "two", 3: "three", 4: "four", 5: "five", 6: "six",
              7: "seven", 8: "eight", 9: "nine", 10: "ten"}

_PRED_PHRASE = {"left_of": "to the left of", "right_of": "to the right of",
                "above": "above", "below": "below", "inside": "inside",
                "in_front_of": "in front of", "behind": "behind"}

_STRATUM_ORDER = {s: i for i, s in enumerate(STRATUM_PRECEDENCE + ("action",))}


def _num(n: int) -> str:
    return _NUM_WORDS.get(n, str(n))


def _plural(label: str, n: int) -> str:
    if n == 1:
        return label
    if label == "person":
        return "people"
    return label if label.endswith("s") else label + "s"


def _ref(entity_id: str) -> str:
    """person_1 -> "person 1"; dog -> "the dog"."""
    if "_" in entity_id and entity_id.rsplit("_", 1)[1].isdigit():
        base, num = entity_id.rsplit("_", 1)
        return f"{base.replace('_', ' ')} {num}"
    return f"the {entity_id.replace('_', ' ')}"


def _cap(sentence: str) -> str:
    return sentence[0].upper() + sentence[1:] if sentence else sentence


def _join(parts: list) -> str:
    parts = [p for p in parts if p]
    if not parts:
        return ""
    if len(parts) == 1:
        return parts[0]
    return ", ".join(parts[:-1]) + " and " + parts[-1]


def _ordered_attrs(entity: dict, max_attrs: int) -> list:
    attrs = sorted(entity.get("attributes", []),
                   key=lambda a: (-a["consensus"], -a["ownership"]["confidence"],
                                  _STRATUM_ORDER.get(a["stratum"], 99), a["text"]))
    return attrs[:max_attrs]


def _entity_clause(entity: dict, intro: str, max_attrs: int, n_entities: int) -> str:
    bits = [f"{intro} in the {entity['position']['grid']} of the frame"]
    d = entity.get("depth")
    if d and n_entities > 1:
        if d["rank"] == 1:
            bits.append("nearest to the camera")
        elif d["rank"] == n_entities:
            bits.append("farthest from the camera")
    attrs = _ordered_attrs(entity, max_attrs)
    # pose/action attributes read as participles ("…, sitting"), not "with sitting"
    plain = [a["text"] for a in attrs if a["stratum"] not in ("action", "pose")]
    acts = [a["text"] for a in attrs if a["stratum"] in ("action", "pose")]
    s = ", ".join(bits)
    if plain:
        s += f", with {_join(plain)}"
    if acts:
        s += f", {_join(acts)}"
    return _cap(s + ".")


def fused_prompt(scene: dict, max_attrs_per_entity: int = 6, max_relations: int = 6,
                 max_basin: int = 4) -> str:
    sentences = []

    # 1) counts
    by_label = scene.get("counts", {}).get("by_label", {})
    if by_label:
        parts = [f"{_num(n)} {_plural(lab, n)}"
                 for lab, n in sorted(by_label.items(), key=lambda kv: (-kv[1], kv[0]))]
        sentences.append(_cap(_join(parts) + "."))

    # 2) entities, primary first, then by saliency rank
    entities = scene.get("entities", [])
    ordered = sorted(entities, key=lambda e: e["saliency"]["rank"])
    n_ent = len(entities)
    for k, e in enumerate(ordered):
        if k == 0:
            intro = f"the primary subject is a {e['label']}" if n_ent > 1 else f"a {e['label']}"
        else:
            intro = f"{_ref(e['id'])} is"
        sentences.append(_entity_clause(e, intro, max_attrs_per_entity, n_ent))

    # 3) relations (strongest-confidence first, capped)
    rels = sorted(scene.get("relations", []),
                  key=lambda r: (-r["confidence"], r["a"], r["b"]))[:max_relations]
    for r in rels:
        phrases = [_PRED_PHRASE[p] for p in r.get("predicates", []) if p in _PRED_PHRASE]
        if phrases:
            sentences.append(_cap(f"{_ref(r['a'])} is {_join(phrases)} {_ref(r['b'])}."))

    # 4) shared basin — uncertainty rendered, never guessed
    for b in scene.get("shared_basin", [])[:max_basin]:
        cands = [_ref(c["entity_id"]) for c in b.get("candidates", [])[:3]]
        if cands:
            joined = cands[0] if len(cands) == 1 else " or ".join(cands)
            sentences.append(_cap(f"one of them ({joined}) has {b['text']}."))
        else:
            sentences.append(_cap(f"somewhere in the scene: {b['text']}."))

    # 5) scene
    sc = scene.get("scene", {})
    bits = []
    setting = (sc.get("setting") or {}).get("value")
    if setting and setting != "unknown":
        bits.append(f"{setting} scene")
    style = (sc.get("style") or {}).get("value")
    if style and style not in ("other", "unknown"):
        bits.append(f"{style} style")
    mood = (sc.get("mood") or {}).get("value")
    if mood:
        bits.append(f"{mood} mood")
    layout = sc.get("layout")
    if layout and layout not in ("unknown", "scattered"):
        bits.append(f"{layout.replace('_', ' ')} composition")
    if (sc.get("symmetry") or {}).get("axis", "none") != "none":
        bits.append(f"{sc['symmetry']['axis']} symmetry")
    if bits:
        sentences.append(_cap(", ".join(bits) + "."))
    for act in sc.get("actions", [])[:3]:
        sentences.append(_cap(f"action in the scene: {act['text']}."))
    for sa in sc.get("scene_attributes", [])[:4]:
        sentences.append(_cap(f"{sa['text']}."))
    ocr_text = (sc.get("ocr") or {}).get("full_text", "").strip()
    if ocr_text:
        sentences.append(_cap(f'visible text: "{ocr_text[:120]}".'))

    return " ".join(sentences)