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
#!/usr/bin/env python3
# SPDX-License-Identifier: MIT OR Apache-2.0
"""Find the best inject feature for a target word in a raw vocab-scan JSON.
The default `pick_features.py` picks suppress features by max `cosine` of
their #1 token within a rime cluster. That's the right metric for the
suppress side (we want features that broadly cover the rime).
For the inject side, we want a feature whose decoder vector specifically
encodes the *target inject word's* embedding direction — not a feature
whose #1 token happens to be a different member of the contrast cluster.
This script scans the raw `vocab_scan_*.json` (gitignored, ~500 MB to
1.5 GB) and, for each requested target word, ranks all features by the
cosine they assign to that word in their top-K. The script reports
top-5 candidates per target.
"""
import json
import sys
from collections import defaultdict
def normalize(text):
"""Match the filter's normalisation: strip leading underscore + space,
case-fold. Returns the cleaned form for word equality testing."""
t = text
# Common BPE leading-space marker variants.
while t and t[0] in (" ", "_", "▁"):
t = t[1:]
return t.casefold()
def find_for_target(features, target):
"""Return list of (max_cosine_for_target, feature_dict) sorted desc."""
target_norm = normalize(target)
hits = []
for ft in features:
best = None
for tok in ft["top_tokens"]:
if normalize(tok["text"]) == target_norm:
cos = tok["cosine"]
if best is None or cos > best:
best = cos
if best is not None:
hits.append((best, ft))
hits.sort(key=lambda x: -x[0])
return hits
def show(raw_path, targets, topn=5):
with open(raw_path, encoding="utf-8") as f:
data = json.load(f)
feats = data["features"]
print(f"=== {raw_path} ({len(feats)} features scanned) ===")
for target in targets:
hits = find_for_target(feats, target)
print(f"\n Top {topn} features for inject_word \"{target}\":")
if not hits:
print(f" (no feature has \"{target}\" in its top-{data.get('top_k','?')})")
continue
for cos, ft in hits[:topn]:
top5 = ", ".join(
t["text"].strip() + f"({t['cosine']:.3f})"
for t in ft["top_tokens"][:5]
)
print(
f" L{ft['feature']['layer']:>2}:{ft['feature']['index']:>5} "
f"cos_to_target={cos:.4f} max_cosine={ft['max_cosine']:.4f} "
f"[{top5}]"
)
def main():
sys.stdout.reconfigure(encoding="utf-8")
if len(sys.argv) < 3:
print(
"Usage: pick_inject_feature.py <raw_scan_json> "
"<target_word> [<target_word> ...]"
)
sys.exit(1)
raw_path = sys.argv[1]
targets = sys.argv[2:]
show(raw_path, targets)
if __name__ == "__main__":
main()