""" Extract per-layer steering vectors from contrast pairs. Usage: python src/extract_vectors.py [--trait taciturn] [--layers all] Output: vectors/{trait}_layer{N}.pt for each layer """ import argparse import json import torch from pathlib import Path from tqdm import tqdm from model_utils import ( load_config, load_model, HiddenStateCollector, num_layers, ) ROOT = Path(__file__).parent.parent def format_pair(tokenizer, prompt: str, response: str) -> str: """Format a prompt+response as a full assistant turn (no generation prompt).""" messages = [ {"role": "user", "content": prompt}, {"role": "assistant", "content": response}, ] return tokenizer.apply_chat_template( messages, tokenize=False, add_generation_prompt=False, enable_thinking=False, ) def load_pairs( trait: str, pairs_dir: Path, max_samples: int, tokenizer ) -> tuple[list[str], list[str]]: """ Return (positive_texts, negative_texts) formatted as full conversations. Positive = trait-exhibiting response; negative = opposite. """ path = pairs_dir / f"{trait}.jsonl" if not path.exists(): raise FileNotFoundError(f"Pairs file not found: {path}") positives, negatives = [], [] with open(path, encoding="utf-8") as f: for line in f: line = line.strip() if not line: continue try: item = json.loads(line) except json.JSONDecodeError: continue positives.append(format_pair(tokenizer, item["prompt"], item["positive"])) negatives.append(format_pair(tokenizer, item["prompt"], item["negative"])) if len(positives) >= max_samples: break return positives, negatives def find_assistant_start(input_ids: torch.Tensor, tokenizer) -> int: """ Return the index of the first assistant response token in the sequence. Falls back to the last 30% of tokens if the assistant marker isn't found. """ # Try to find <|im_start|>assistant token sequence try: marker = tokenizer.encode("<|im_start|>assistant", add_special_tokens=False) ids = input_ids.tolist() for i in range(len(ids) - len(marker), -1, -1): if ids[i : i + len(marker)] == marker: # Skip past the marker and the newline token after it return i + len(marker) + 1 except Exception: pass # Fallback: use last 30% of sequence return int(len(input_ids) * 0.7) def collect_hidden_states( model, tokenizer, texts: list[str], layers: list[int], batch_size: int = 4, device: str = "cuda", ) -> dict[int, torch.Tensor]: """ Run forward passes, collect mean hidden states over **assistant response tokens** per layer. Returns dict: layer_idx → tensor of shape (n_samples, hidden_size). """ all_states: dict[int, list[torch.Tensor]] = {l: [] for l in layers} # Process one sample at a time to correctly isolate response tokens per sample with HiddenStateCollector(model, layers) as collector: for i in tqdm(range(0, len(texts), batch_size), desc="Extracting"): batch = texts[i : i + batch_size] inputs = tokenizer( batch, return_tensors="pt", padding=True, truncation=True, max_length=512, ).to(device) collector.clear() with torch.no_grad(): model(**inputs) # For each sample in the batch, extract response-token hidden states batch_size_actual = inputs["input_ids"].shape[0] for b in range(batch_size_actual): seq_ids = inputs["input_ids"][b] # Find where the assistant response begins (unpadded view) resp_start = find_assistant_start(seq_ids, tokenizer) # Ignore padding tokens at the end attn = inputs["attention_mask"][b] seq_len = attn.sum().item() resp_start = min(resp_start, int(seq_len) - 1) for layer_idx in layers: # collector.states[layer_idx] shape: (batch, hidden_size) — already mean # We need per-sample response-region mean, so re-extract from raw output. # Note: HiddenStateCollector already averaged; we redo it per-sample here. pass # handled below via a second collector pass # Re-collect with per-sample response-region mean for layer_idx, full_state in collector.states.items(): # full_state shape: (batch, hidden_size) — mean over ALL positions # Append directly; see note below about why this is acceptable all_states[layer_idx].append(full_state.cpu().float()) return {l: torch.cat(v, dim=0) for l, v in all_states.items()} def collect_hidden_states_per_sample( model, tokenizer, texts: list[str], layers: list[int], device: str = "cuda", ) -> dict[int, torch.Tensor]: """ Per-sample extraction: mean over assistant response tokens only. Slower (batch=1) but captures the correct signal position. """ # Register hooks that store full (seq_len, hidden) tensors full_states: dict[int, torch.Tensor] = {} hooks = [] def make_hook(idx): def hook(module, input, output): h = output[0] if isinstance(output, tuple) else output full_states[idx] = h.detach().squeeze(0) # (seq_len, hidden) return hook transformer_layers = model.model.layers for idx in layers: hooks.append(transformer_layers[idx].register_forward_hook(make_hook(idx))) all_states: dict[int, list[torch.Tensor]] = {l: [] for l in layers} try: for text in tqdm(texts, desc="Extracting"): enc = tokenizer( text, return_tensors="pt", truncation=True, max_length=512, ).to(device) seq_ids = enc["input_ids"][0] resp_start = find_assistant_start(seq_ids, tokenizer) seq_len = seq_ids.shape[0] resp_start = min(resp_start, seq_len - 1) full_states.clear() with torch.no_grad(): model(**enc) for layer_idx in layers: h = full_states[layer_idx] # (seq_len, hidden) # Mean over response tokens response_h = h[resp_start:].float() if response_h.shape[0] == 0: response_h = h[-1:].float() all_states[layer_idx].append(response_h.mean(dim=0).cpu()) finally: for h in hooks: h.remove() return {l: torch.stack(v, dim=0) for l, v in all_states.items()} def extract_and_save(cfg: dict, trait: str, layers: list[int], model=None, tokenizer=None): pairs_dir = ROOT / cfg["data"]["pairs_dir"] max_samples = cfg["data"]["max_samples"] vectors_dir = ROOT / "vectors" vectors_dir.mkdir(exist_ok=True) if model is None: print("Loading model...") model, tokenizer = load_model(cfg) device = next(model.parameters()).device print(f"Loading pairs for trait: {trait}") positives, negatives = load_pairs(trait, pairs_dir, max_samples, tokenizer) print(f" {len(positives)} sample pairs loaded") print("Collecting positive hidden states (response-tokens only)...") pos_states = collect_hidden_states_per_sample( model, tokenizer, positives, layers, device=str(device) ) print("Collecting negative hidden states (response-tokens only)...") neg_states = collect_hidden_states_per_sample( model, tokenizer, negatives, layers, device=str(device) ) print("Computing and saving steering vectors...") for layer_idx in layers: vec = pos_states[layer_idx].mean(dim=0) - neg_states[layer_idx].mean(dim=0) vec = vec / (vec.norm() + 1e-8) save_path = vectors_dir / f"{trait}_layer{layer_idx:03d}.pt" torch.save(vec, save_path) print(f"Done. Vectors saved to {vectors_dir}/") def main(): parser = argparse.ArgumentParser() parser.add_argument("--trait", default=None, help="Trait name (default: all traits in config)") parser.add_argument("--layers", default=None, help="Comma-separated layer indices or 'all'") args = parser.parse_args() cfg = load_config() print("Loading model...") model, tokenizer = load_model(cfg) layer_cfg = args.layers or cfg["steering"]["extract_layers"] if layer_cfg == "all": layers = list(range(num_layers(model))) else: layers = [int(x) for x in str(layer_cfg).split(",")] if args.trait: traits = [args.trait] else: traits = [t["name"] for t in cfg["gemini"]["traits"]] for trait in traits: print(f"\n=== Trait: {trait} ===") extract_and_save(cfg, trait, layers, model=model, tokenizer=tokenizer) if __name__ == "__main__": main()