cil/chart_features.py

from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

try:
    import numpy as np
except ImportError:  # pragma: no cover
    np = None


CONTEXT_HASH_WIDTH = 8
OBSERVATION_EMBED_DIM = 32
OBJECT_LAYOUT_EMBED_DIM = 64
CHART_FEATURE_MODES = (
    "base",
    "base_context",
    "base_context_obs",
    "base_context_obj",
    "base_context_obs_obj",
)
_EMBEDDING_CACHE: dict[str, Any] = {}


def build_chart_feature(
    base_action: Any,
    metadata: dict[str, Any] | None = None,
    *,
    mode: str = "base",
) -> Any:
    """Build a deployment-visible chart feature vector.

    `base` preserves the original behavior: the chart token is the flattened
    base action chunk. `base_context` appends stable hashes of task/language
    metadata that are visible at proposal time. `base_context_obs` additionally
    appends a precomputed observation embedding referenced by metadata.
    `base_context_obj` appends a precomputed RGB object-layout embedding, and
    `base_context_obs_obj` appends both. These modes intentionally do not read
    outcomes, labels, hidden chart branches, or evaluator-only fields.
    """

    if np is None:  # pragma: no cover
        raise ImportError("build_chart_feature requires numpy")
    if mode not in CHART_FEATURE_MODES:
        raise ValueError(f"unknown chart feature mode: {mode}")
    base = np.asarray(base_action, dtype=np.float32).reshape(-1)
    if mode == "base":
        return base
    metadata = metadata or {}
    context = np.asarray(
        [
            *_stable_hash_features(str(metadata.get("task_id", "")), CONTEXT_HASH_WIDTH),
            *_stable_hash_features(str(metadata.get("instruction", "")).lower(), CONTEXT_HASH_WIDTH),
            min(len(str(metadata.get("instruction", ""))) / 128.0, 4.0),
            min(len(str(metadata.get("instruction", "")).split()) / 32.0, 4.0),
        ],
        dtype=np.float32,
    )
    if mode in {"base_context_obs", "base_context_obs_obj"}:
        obs = _load_embedding(
            metadata.get("observation_embedding_path"),
            dim=OBSERVATION_EMBED_DIM,
            chart_root=metadata.get("_chart_root"),
        )
        context = np.concatenate([context, obs.astype(np.float32, copy=False)])
    if mode in {"base_context_obj", "base_context_obs_obj"}:
        obj = _load_embedding(
            metadata.get("object_embedding_path"),
            dim=OBJECT_LAYOUT_EMBED_DIM,
            chart_root=metadata.get("_chart_root"),
        )
        context = np.concatenate([context, obj.astype(np.float32, copy=False)])
    return np.concatenate([base, context]).astype(np.float32, copy=False)


def chart_feature_dim(base_action: Any, *, mode: str = "base") -> int:
    return int(build_chart_feature(base_action, {}, mode=mode).reshape(-1).shape[0])


def _stable_hash_features(text: str, width: int) -> list[float]:
    digest = hashlib.sha256(text.encode("utf-8")).digest()
    return [float(digest[index] / 127.5 - 1.0) for index in range(width)]


def _load_embedding(value: Any, *, dim: int, chart_root: Any = None) -> Any:
    if np is None:  # pragma: no cover
        raise ImportError("build_chart_feature requires numpy")
    if not value:
        return np.zeros(dim, dtype=np.float32)
    path_text, dataset, row_index = _parse_embedding_ref(str(value))
    path = Path(path_text)
    if not path.is_absolute() and chart_root:
        path = Path(str(chart_root)) / path
    cache_key = str(path.resolve())
    if cache_key not in _EMBEDDING_CACHE:
        with np.load(path, allow_pickle=False) as data:
            _EMBEDDING_CACHE[cache_key] = np.asarray(data[dataset], dtype=np.float32)
    matrix = _EMBEDDING_CACHE[cache_key]
    vector = np.asarray(matrix[int(row_index)], dtype=np.float32).reshape(-1)
    if vector.shape[0] == dim:
        return vector
    output = np.zeros(dim, dtype=np.float32)
    width = min(dim, vector.shape[0])
    output[:width] = vector[:width]
    return output


def _parse_embedding_ref(value: str) -> tuple[str, str, int]:
    if "#" not in value:
        return value, "embeddings", 0
    path_text, ref = value.split("#", 1)
    parts = [part for part in ref.split("/") if part]
    if len(parts) != 2:
        raise ValueError(f"invalid observation embedding ref: {value}")
    return path_text, parts[0], int(parts[1])