| """Extract the pain direction from contrastive pairs and save it as safetensors. |
| |
| python scripts/extract_pain_vector.py \ |
| --model Qwen/Qwen3-8B --pairs data/contrastive_pairs.jsonl \ |
| --layers mid --out vectors/pain_qwen3-8b.safetensors |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import time |
|
|
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) |
|
|
| from codelock_affect import model as cm |
| from codelock_affect.steering import ( |
| decoder_layers, |
| extract_vector, |
| parse_layers, |
| ) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model", default=cm.DEFAULT_MODEL) |
| ap.add_argument("--pairs", default="data/contrastive_pairs.jsonl") |
| ap.add_argument("--layers", default="mid", help="'mid', '10-26' or '8,12,16'") |
| ap.add_argument("--out", default="vectors/pain.safetensors") |
| ap.add_argument("--four-bit", action="store_true") |
| ap.add_argument( |
| "--model-id", |
| default=None, |
| help="hub id of the base model recorded in the vector metadata (defaults to --model)", |
| ) |
| args = ap.parse_args() |
|
|
| with open(args.pairs, encoding="utf-8") as f: |
| pairs = [json.loads(line) for line in f if line.strip()] |
| print(f"loading {args.model} ...") |
| model, tok = cm.load(args.model, four_bit=args.four_bit) |
| layers = parse_layers(args.layers, len(decoder_layers(model))) |
| print(f"{len(pairs)} pairs, layers {layers[0]}-{layers[-1]}") |
|
|
| t0 = time.time() |
| vec = extract_vector(model, tok, pairs, layers, name="pain", log=print) |
| vec.meta["base_model"] = args.model_id or args.model |
| os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True) |
| vec.save(args.out) |
| print(f"saved {args.out} in {time.time() - t0:.1f}s") |
| for i in layers: |
| print(f" layer {i:2d} |h| = {vec.norms[i]:8.2f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
