chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
Executable
+229
@@ -0,0 +1,229 @@
|
||||
#!/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()
|
||||
Executable
+64
@@ -0,0 +1,64 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Run a small steering scale sweep through ds4.
|
||||
|
||||
This is intentionally thin: it exercises the same public CLI options users
|
||||
will use in production and leaves all inference behavior inside ds4.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
def read_prompts(path: Path) -> list[str]:
|
||||
prompts = []
|
||||
for line in path.read_text(encoding="utf-8").splitlines():
|
||||
line = line.strip()
|
||||
if line and not line.startswith("#"):
|
||||
prompts.append(line)
|
||||
if not prompts:
|
||||
raise SystemExit(f"{path}: no prompts found")
|
||||
return prompts
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--ds4", default="./ds4")
|
||||
ap.add_argument("--model", default="ds4flash.gguf")
|
||||
ap.add_argument("--direction", required=True,
|
||||
help="flat f32 vector file produced by build_direction.py")
|
||||
ap.add_argument("--prompts", required=True)
|
||||
ap.add_argument("--scales", default="-2,-1,-0.5,0,0.5,1,2")
|
||||
ap.add_argument("--tokens", type=int, default=160)
|
||||
ap.add_argument("--ctx", type=int, default=4096)
|
||||
ap.add_argument("--attn-scale", type=float, default=0.0)
|
||||
ap.add_argument("--nothink", action="store_true")
|
||||
args = ap.parse_args()
|
||||
|
||||
prompts = read_prompts(Path(args.prompts))
|
||||
scales = [float(x) for x in args.scales.split(",") if x.strip()]
|
||||
|
||||
for prompt in prompts:
|
||||
print("=" * 80)
|
||||
print(f"PROMPT: {prompt}")
|
||||
for scale in scales:
|
||||
print("-" * 80)
|
||||
print(f"FFN scale: {scale:g}")
|
||||
cmd = [
|
||||
args.ds4,
|
||||
"-m", args.model,
|
||||
"--ctx", str(args.ctx),
|
||||
"-n", str(args.tokens),
|
||||
"--temp", "0",
|
||||
"--dir-steering-file", args.direction,
|
||||
"--dir-steering-ffn", str(scale),
|
||||
"--dir-steering-attn", str(args.attn_scale),
|
||||
"-p", prompt,
|
||||
]
|
||||
if args.nothink:
|
||||
cmd.append("--nothink")
|
||||
subprocess.run(cmd, check=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user