wehub-resource-sync
230 lines
7.7 KiB
Python
Executable File
230 lines
7.7 KiB
Python
Executable File
#!/usr/bin/env python3
|
||
"""Build a DS4 directional-steering vector from paired prompt sets.
|
||
|
||
The extractor asks ds4 to dump one 4096-wide activation row per layer, averages
|
||
the target and control rows, and writes a flat f32 file with 43 layer vectors.
|
||
At runtime ds4 applies:
|
||
|
||
y = y - scale * direction[layer] * dot(direction[layer], y)
|
||
|
||
Positive scale suppresses the target direction. Negative scale amplifies it.
|
||
"""
|
||
|
||
import argparse
|
||
import array
|
||
import json
|
||
import math
|
||
import os
|
||
import subprocess
|
||
import tempfile
|
||
from pathlib import Path
|
||
|
||
|
||
N_LAYER = 43
|
||
N_EMBD = 4096
|
||
|
||
SPECIALS = {
|
||
"bos": "<｜begin▁of▁sentence｜>",
|
||
"user": "<｜User｜>",
|
||
"assistant": "<｜Assistant｜>",
|
||
"think": "<think>",
|
||
"nothink": "</think>",
|
||
}
|
||
|
||
|
||
def read_prompt_file(path: Path) -> list[str]:
|
||
"""Read one prompt per non-empty line, ignoring shell-style comments."""
|
||
prompts: list[str] = []
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
line = line.strip()
|
||
if not line or line.startswith("#"):
|
||
continue
|
||
prompts.append(line)
|
||
if not prompts:
|
||
raise SystemExit(f"{path}: no prompts found")
|
||
return prompts
|
||
|
||
|
||
def render_ds4_prompt(system: str, user: str, think: bool) -> str:
|
||
"""Render the minimal DS4 chat prefix used for activation capture."""
|
||
pieces = [SPECIALS["bos"]]
|
||
if system:
|
||
pieces.append(system)
|
||
pieces += [
|
||
SPECIALS["user"],
|
||
user,
|
||
SPECIALS["assistant"],
|
||
SPECIALS["think"] if think else SPECIALS["nothink"],
|
||
]
|
||
return "".join(pieces)
|
||
|
||
|
||
def normalize(v: list[float]) -> list[float]:
|
||
n2 = sum(x * x for x in v)
|
||
if n2 <= 0.0:
|
||
return v
|
||
inv = 1.0 / math.sqrt(n2)
|
||
return [x * inv for x in v]
|
||
|
||
|
||
def dot(a: list[float], b: list[float]) -> float:
|
||
return sum(x * y for x, y in zip(a, b))
|
||
|
||
|
||
def run_capture(
|
||
ds4: Path,
|
||
model: Path,
|
||
prompt: str,
|
||
system: str,
|
||
think: bool,
|
||
ctx: int,
|
||
component: str,
|
||
work: Path,
|
||
) -> list[list[float]]:
|
||
"""Run ds4 once and return the last prompt-row dump for every layer."""
|
||
prompt_path = work / "prompt.txt"
|
||
prompt_path.write_text(render_ds4_prompt(system, prompt, think), encoding="utf-8")
|
||
dump_prefix = work / "dump"
|
||
|
||
env = os.environ.copy()
|
||
env["DS4_METAL_GRAPH_DUMP_PREFIX"] = str(dump_prefix)
|
||
env["DS4_METAL_GRAPH_DUMP_NAME"] = component
|
||
env["DS4_METAL_GRAPH_DUMP_POS"] = "0"
|
||
|
||
cmd = [
|
||
str(ds4),
|
||
"-m", str(model),
|
||
"--ctx", str(ctx),
|
||
"--prompt-file", str(prompt_path),
|
||
"-n", "1",
|
||
]
|
||
subprocess.run(cmd, cwd=ds4.parent, env=env, check=True,
|
||
stdout=subprocess.DEVNULL, stderr=subprocess.PIPE)
|
||
|
||
rows: list[list[float]] = []
|
||
for layer in range(N_LAYER):
|
||
path = work / f"dump_{component}-{layer}_pos0.bin"
|
||
data = array.array("f")
|
||
with path.open("rb") as f:
|
||
data.fromfile(f, path.stat().st_size // 4)
|
||
if len(data) < N_EMBD or len(data) % N_EMBD != 0:
|
||
raise RuntimeError(f"bad dump shape for {path}: {len(data)} floats")
|
||
rows.append(list(data[-N_EMBD:]))
|
||
return rows
|
||
|
||
|
||
def add_rows(total: list[list[float]], rows: list[list[float]]) -> None:
|
||
for layer in range(N_LAYER):
|
||
dst = total[layer]
|
||
src = rows[layer]
|
||
for i, value in enumerate(src):
|
||
dst[i] += value
|
||
|
||
|
||
def main() -> None:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--ds4", default="./ds4", help="path to the ds4 CLI")
|
||
ap.add_argument("--model", default="ds4flash.gguf", help="GGUF model path")
|
||
ap.add_argument("--good-file", required=True,
|
||
help="desired/target prompts, one per line")
|
||
ap.add_argument("--bad-file", required=True,
|
||
help="contrast/control prompts, one per line")
|
||
ap.add_argument("--out", default="dir-steering/out/direction.json",
|
||
help="metadata JSON path; .f32 is written next to it")
|
||
ap.add_argument("--ctx", type=int, default=512)
|
||
ap.add_argument("--system", default="You are a helpful assistant.")
|
||
ap.add_argument("--component", default="ffn_out",
|
||
choices=("ffn_out", "attn_out"),
|
||
help="runtime-editable 4096-wide activation stream")
|
||
ap.add_argument("--think", action="store_true",
|
||
help="capture after <think>; default captures direct answers")
|
||
ap.add_argument("--pair-normalize", action="store_true",
|
||
help="average normalized per-pair differences")
|
||
ap.add_argument("--no-orthogonalize", action="store_true",
|
||
help="do not remove the component parallel to the control mean")
|
||
args = ap.parse_args()
|
||
|
||
ds4 = Path(args.ds4).resolve()
|
||
model = Path(args.model).resolve()
|
||
good_prompts = read_prompt_file(Path(args.good_file))
|
||
bad_prompts = read_prompt_file(Path(args.bad_file))
|
||
n = min(len(good_prompts), len(bad_prompts))
|
||
good_prompts = good_prompts[:n]
|
||
bad_prompts = bad_prompts[:n]
|
||
|
||
good_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||
bad_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||
pair_sum = [[0.0] * N_EMBD for _ in range(N_LAYER)]
|
||
|
||
with tempfile.TemporaryDirectory(prefix="ds4-dir-steer-") as td:
|
||
root = Path(td)
|
||
for i, (good, bad) in enumerate(zip(good_prompts, bad_prompts), 1):
|
||
print(f"pair {i}/{n}", flush=True)
|
||
gw = root / f"good-{i}"
|
||
bw = root / f"bad-{i}"
|
||
gw.mkdir()
|
||
bw.mkdir()
|
||
good_rows = run_capture(ds4, model, good, args.system, args.think,
|
||
args.ctx, args.component, gw)
|
||
bad_rows = run_capture(ds4, model, bad, args.system, args.think,
|
||
args.ctx, args.component, bw)
|
||
add_rows(good_sum, good_rows)
|
||
add_rows(bad_sum, bad_rows)
|
||
if args.pair_normalize:
|
||
for layer in range(N_LAYER):
|
||
diff = normalize([
|
||
good_rows[layer][j] - bad_rows[layer][j]
|
||
for j in range(N_EMBD)
|
||
])
|
||
for j, value in enumerate(diff):
|
||
pair_sum[layer][j] += value
|
||
|
||
layers = []
|
||
for layer in range(N_LAYER):
|
||
good_mean = [x / n for x in good_sum[layer]]
|
||
bad_mean = [x / n for x in bad_sum[layer]]
|
||
if args.pair_normalize:
|
||
direction = normalize([x / n for x in pair_sum[layer]])
|
||
else:
|
||
direction = normalize([
|
||
good_mean[i] - bad_mean[i]
|
||
for i in range(N_EMBD)
|
||
])
|
||
if not args.no_orthogonalize:
|
||
base = normalize(bad_mean)
|
||
projection = dot(direction, base)
|
||
direction = normalize([
|
||
direction[i] - projection * base[i]
|
||
for i in range(N_EMBD)
|
||
])
|
||
layers.append(direction)
|
||
|
||
out = Path(args.out)
|
||
out.parent.mkdir(parents=True, exist_ok=True)
|
||
payload = {
|
||
"format": "ds4-directional-steering-v1",
|
||
"shape": [N_LAYER, N_EMBD],
|
||
"component": args.component,
|
||
"thinking": bool(args.think),
|
||
"pair_normalize": bool(args.pair_normalize),
|
||
"orthogonalize_control_mean": not args.no_orthogonalize,
|
||
"good_file": str(Path(args.good_file)),
|
||
"bad_file": str(Path(args.bad_file)),
|
||
"model": str(model),
|
||
"note": "runtime positive scale suppresses this direction; negative scale amplifies it",
|
||
}
|
||
out.write_text(json.dumps(payload, indent=2), encoding="utf-8")
|
||
|
||
flat = array.array("f")
|
||
for direction in layers:
|
||
flat.extend(direction)
|
||
f32_out = out.with_suffix(".f32")
|
||
with f32_out.open("wb") as f:
|
||
flat.tofile(f)
|
||
print(f"wrote {out}")
|
||
print(f"wrote {f32_out}")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|