| """Calibrate four residual directions using only the grouped training split.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| from codelock_affect.channels import CHANNELS |
| from codelock_affect.config import ProjectConfig |
| from codelock_affect.extraction import extract_vector |
| from codelock_affect.layers import decoder_layers, parse_layers |
| from codelock_affect.model import load |
| from codelock_affect.vectors import VectorBank |
| from codelock_affect.evaluation import environment, sha256_file, write_json |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--config", default="configs/default.json") |
| p.add_argument("--model") |
| p.add_argument("--revision") |
| p.add_argument("--data", default="data/affect_pairs.jsonl") |
| p.add_argument("--output", default="vectors/multichannel") |
| p.add_argument("--four-bit", action="store_true") |
| args = p.parse_args() |
| cfg = ProjectConfig.load(args.config) |
| pairs = [ |
| json.loads(line) |
| for line in Path(args.data).read_text(encoding="utf-8").splitlines() |
| if line.strip() |
| ] |
| model, tok = load(args.model or cfg.model, args.four_bit, revision=args.revision) |
| layers = parse_layers(cfg.layers, len(decoder_layers(model))) |
| vectors = {} |
| for channel in CHANNELS: |
| examples = [ |
| row |
| for row in pairs |
| if row["channel"] == channel and row["split"] == "train" |
| ] |
| vector = extract_vector( |
| model, tok, examples, layers, channel, lambda text: print(text, flush=True) |
| ) |
| vector.meta.update( |
| { |
| "dataset_file_sha256": sha256_file(args.data), |
| "quantized_4bit": args.four_bit, |
| "environment": environment(), |
| "requested_revision": args.revision, |
| } |
| ) |
| vectors[channel] = vector |
| |
| vector.save(Path(args.output) / f"{channel}.safetensors") |
| bank = VectorBank(vectors) |
| bank.save(args.output) |
| write_json( |
| Path(args.output) / "extraction.json", |
| { |
| "environment": environment(), |
| "config": cfg.to_dict(), |
| "dataset_sha256": sha256_file(args.data), |
| "channels": {c: v.meta for c, v in vectors.items()}, |
| }, |
| ) |
| print(f"Saved {len(vectors)} channels to {args.output}", flush=True) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
