add context learned dominance diagnostics

workspace/scripts/eval_learned_dominance_selector.py

@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import argparse
+import hashlib
 import json
 import math
 import re
@@ -34,6 +35,7 @@ BASIC_FEATURE_NAMES = [
     "num_candidates",
 ]
 FEATURE_NAMES = BASIC_FEATURE_NAMES
+CONTEXT_HASH_WIDTH = 8
 
 
 def main(argv: list[str] | None = None) -> int:
@@ -56,7 +58,7 @@ def main(argv: list[str] | None = None) -> int:
     parser.add_argument("--ridge-lambdas", default="0,0.01,0.1,1,10,100")
     parser.add_argument(
         "--feature-set",
-        choices=("basic", "tangent"),
+        choices=("basic", "tangent", "context", "context_tangent"),
         default="basic",
         help="Deployment-visible feature family for candidate-level dominance fitting.",
     )
@@ -218,6 +220,7 @@ def _candidate_dataset(
         successes = [float(bool(value)) for value in row.get("candidate_success", [])[:k]]
         tangents = row.get("generated_tangents", [])[:k]
         candidate_types = row.get("candidate_types", [])[:k]
+        source_task_ids = row.get("candidate_source_task_ids", [])[:k]
         if not scores or len(scores) != len(utilities):
             continue
         score_mean = sum(scores) / len(scores)
@@ -234,6 +237,13 @@ def _candidate_dataset(
                 score_std=score_std,
                 candidate_index=candidate_index,
                 candidate_type=candidate_types[candidate_index] if candidate_index < len(candidate_types) else "",
+                context={
+                    "target_task_id": row.get("task_id", ""),
+                    "instruction": row.get("instruction", ""),
+                    "source_task_id": source_task_ids[candidate_index]
+                    if candidate_index < len(source_task_ids)
+                    else "",
+                },
                 tangent=tangent,
                 num_candidates=len(scores),
                 feature_set=feature_set,
@@ -276,12 +286,24 @@ def _candidate_dataset(
 def _feature_names(feature_set: str) -> list[str]:
     if feature_set == "basic":
         return list(BASIC_FEATURE_NAMES)
+    context_names = [
+        *[f"target_task_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)],
+        *[f"source_task_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)],
+        *[f"instruction_hash_{index:02d}" for index in range(CONTEXT_HASH_WIDTH)],
+        "source_target_same_task",
+        "instruction_chars_per_128",
+        "instruction_words_per_32",
+    ]
+    tangent_names = [
+        *[f"tangent_{index:02d}" for index in range(21)],
+        *[f"abs_tangent_{index:02d}" for index in range(21)],
+    ]
     if feature_set == "tangent":
-        return [
-            *BASIC_FEATURE_NAMES,
-            *[f"tangent_{index:02d}" for index in range(21)],
-            *[f"abs_tangent_{index:02d}" for index in range(21)],
-        ]
+        return [*BASIC_FEATURE_NAMES, *tangent_names]
+    if feature_set == "context":
+        return [*BASIC_FEATURE_NAMES, *context_names]
+    if feature_set == "context_tangent":
+        return [*BASIC_FEATURE_NAMES, *context_names, *tangent_names]
     raise ValueError(f"unknown feature_set: {feature_set}")
 
 
@@ -296,6 +318,7 @@ def _candidate_feature(
     tangent: np.ndarray,
     num_candidates: int,
     feature_set: str,
+    context: dict[str, Any] | None = None,
 ) -> np.ndarray:
     tangent = np.asarray(tangent, dtype=float).reshape(-1)
     if tangent.size < 21:
@@ -323,9 +346,42 @@ def _candidate_feature(
         return basic
     if feature_set == "tangent":
         return np.concatenate([basic, tangent.astype(float), np.abs(tangent).astype(float)])
+    if feature_set == "context":
+        return np.concatenate([basic, _context_feature(context or {})])
+    if feature_set == "context_tangent":
+        return np.concatenate(
+            [
+                basic,
+                _context_feature(context or {}),
+                tangent.astype(float),
+                np.abs(tangent).astype(float),
+            ]
+        )
     raise ValueError(f"unknown feature_set: {feature_set}")
 
 
+def _context_feature(context: dict[str, Any]) -> np.ndarray:
+    target_task = str(context.get("target_task_id", ""))
+    source_task = str(context.get("source_task_id", ""))
+    instruction = str(context.get("instruction", ""))
+    return np.asarray(
+        [
+            *_stable_hash_features(target_task, CONTEXT_HASH_WIDTH),
+            *_stable_hash_features(source_task, CONTEXT_HASH_WIDTH),
+            *_stable_hash_features(instruction.lower(), CONTEXT_HASH_WIDTH),
+            float(bool(target_task) and target_task == source_task),
+            min(len(instruction) / 128.0, 4.0),
+            min(len(instruction.split()) / 32.0, 4.0),
+        ],
+        dtype=float,
+    )
+
+
+def _stable_hash_features(text: str, width: int) -> list[float]:
+    digest = hashlib.sha256(text.encode("utf-8")).digest()
+    return [((digest[index] / 127.5) - 1.0) for index in range(width)]
+
+
 def _target_value(
     target: str,
     *,