onf-c1/scripts/embed_task_names.py

#!/usr/bin/env python
"""scripts/embed_task_names.py — embed one suite's task instructions into g_lang.npz.

Encodes the graph's task names with the frozen local Qwen2.5-0.5B, centres them on their own mean
and writes them beside that suite's graph artifacts. This is what replaces the string-matching task
lane: see onf.graph.lang for the measured accuracy of the thing being built here, including the
operating points where it is WORSE than the lane it replaces.

The artifact is written outside ArtifactSchemas.NODES, so an existing g_nodes.npz stays valid and a
graph without g_lang.npz simply runs with the language prior switched off.

Two pre-registered gates, both reported and both fatal unless --force:
  separation  the mean cosine between two DIFFERENT tasks must sit below --max-cross. Uncentred,
              Qwen puts every task pair above 0.9 and no temperature can act on that.
  distinctness  every task must be its own nearest neighbour. A graph where two task names collide
              cannot be tilted apart by this mechanism at all.

Usage:
    /home/quang/miniconda3/envs/stablevla/bin/python scripts/embed_task_names.py long
    /home/quang/miniconda3/envs/stablevla/bin/python scripts/embed_task_names.py long --out /tmp/g_lang.npz

Exit codes: 0 = written; 1 = operational error; 2 = a gate failed.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path
from typing import Sequence

import numpy as np

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "src") not in sys.path:
    sys.path.insert(0, str(REPO_ROOT / "src"))

from onf.config import default_paths  # noqa: E402
from onf.graph.core.nodes import NodeTable  # noqa: E402
from onf.graph.lang import (  # noqa: E402
    LANG_NPZ,
    InstructionEncoder,
    TaskEmbeddings,
    normalise_instruction,
)

# Above this mean off-diagonal cosine the embeddings carry no usable contrast. Centred Qwen scores
# well under it on `long`; uncentred it scores ~0.95 and the whole mechanism is inert.
DEFAULT_MAX_CROSS = 0.75


def build_parser() -> argparse.ArgumentParser:
    """The command line.

    Returns:
        The parser.
    """
    p = argparse.ArgumentParser(description="embed a suite's task instructions into g_lang.npz")
    p.add_argument("suite", help="suite name, e.g. long")
    p.add_argument("--graph-dir", default="", help="override the resolved graph directory")
    p.add_argument("--out", default="", help=f"override the output path (default <graph-dir>/{LANG_NPZ})")
    p.add_argument("--model-dir", default="", help="encoder snapshot; default is the local Qwen2.5-0.5B")
    p.add_argument("--max-cross", type=float, default=DEFAULT_MAX_CROSS,
                   help="fail if the mean between-task cosine exceeds this")
    p.add_argument("--force", action="store_true", help="write even if a gate fails")
    return p


def report(emb: TaskEmbeddings, max_cross: float) -> bool:
    """Print the similarity structure and evaluate both gates.

    Args:
        emb: The freshly built embeddings.
        max_cross: Separation gate.

    Returns:
        True when both gates pass.
    """
    sim = emb.e_task @ emb.e_task.T
    off = ~np.eye(emb.n_tasks, dtype=bool)
    cross_mean, cross_max = float(sim[off].mean()), float(sim[off].max())
    nearest = np.argmax(np.where(off, sim, -np.inf), axis=1)

    print(f"encoder    : {emb.model_id}")
    print(f"tasks      : {emb.n_tasks}")
    print(f"cross-task cosine: mean {cross_mean:+.4f}  max {cross_max:+.4f}  (gate < {max_cross})")
    print("\nnearest other task per task (the pairs this prior cannot separate):")
    for i, name in enumerate(emb.task_names):
        j = int(nearest[i])
        print(f"  [{i}] {normalise_instruction(name)[:58]:58s} -> [{j}] cos {sim[i, j]:+.3f}")

    separated = cross_mean < max_cross
    distinct = bool(np.all(np.diag(sim) > sim[np.arange(emb.n_tasks), nearest]))
    if not separated:
        print(f"\nGATE FAILED: mean cross-task cosine {cross_mean:.4f} >= {max_cross}. The task "
              f"embeddings are not separated; check that centering was applied.")
    if not distinct:
        print("\nGATE FAILED: some task is not its own nearest neighbour -- two task names collide.")
    return separated and distinct


def main(argv: Sequence[str] | None = None) -> int:
    """Entry point.

    Args:
        argv: Command line, or None for sys.argv.

    Returns:
        Process exit code.
    """
    args = build_parser().parse_args(argv)
    graph_dir = Path(args.graph_dir) if args.graph_dir else Path(default_paths().graph(args.suite))
    if not graph_dir.is_dir():
        print(f"error: no graph directory at {graph_dir}", file=sys.stderr)
        return 1

    nodes = NodeTable.load(graph_dir / "g_nodes.npz")
    encoder = InstructionEncoder(model_dir=args.model_dir)
    print(f"graph      : {graph_dir}")
    emb = TaskEmbeddings.build(nodes.task_names, encoder)

    passed = report(emb, args.max_cross)
    if not passed and not args.force:
        return 2

    written = emb.save(Path(args.out) if args.out else graph_dir)
    print(f"\nwrote {written}" + ("  (GATE OVERRIDDEN BY --force)" if not passed else ""))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())