File size: 4,351 Bytes
b4269cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | 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])
|