| """Paired answer-token extraction at exactly the decoder outputs we later steer.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from collections.abc import Iterable, Sequence |
| import torch |
| from .layers import decoder_layers, hidden_tensor |
| from .vectors import AffectVector |
|
|
|
|
| def answer_tokens(tokenizer, situation: str, answer: str, maximum: int = 1024): |
| if not answer.strip(): |
| raise ValueError("An empty answer has no extraction region") |
| prompt = tokenizer.apply_chat_template( |
| [{"role": "user", "content": situation}], |
| tokenize=False, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| ) |
| prefix = tokenizer(prompt, add_special_tokens=False)["input_ids"] |
| full = tokenizer(prompt + answer, add_special_tokens=False)["input_ids"] |
| |
| |
| start = 0 |
| while start < min(len(prefix), len(full)) and prefix[start] == full[start]: |
| start += 1 |
| if start >= len(full) or len(full) > maximum: |
| raise ValueError( |
| "No answer tokens or extraction example exceeds maximum length" |
| ) |
| return full, start |
|
|
|
|
| @torch.no_grad() |
| def answer_statistics(model, tokenizer, situation, answer, layers): |
| ids, start = answer_tokens(tokenizer, situation, answer) |
| blocks = decoder_layers(model) |
| captured = {} |
| handles = [] |
|
|
| def hook(layer): |
| def record(_module, _inputs, output): |
| h = hidden_tensor(output)[0, start:].detach().float() |
| captured[layer] = (h.mean(dim=0).cpu(), h.norm(dim=-1).mean().item()) |
|
|
| return record |
|
|
| try: |
| for layer in layers: |
| handles.append(blocks[layer].register_forward_hook(hook(layer))) |
| device = model.get_input_embeddings().weight.device |
| model(input_ids=torch.tensor([ids], device=device), use_cache=False) |
| finally: |
| for handle in handles: |
| handle.remove() |
| if set(captured) != set(layers): |
| raise RuntimeError("Decoder adapter did not observe every requested layer") |
| return captured |
|
|
|
|
| def extract_vector( |
| model, |
| tokenizer, |
| pairs: Iterable[dict], |
| layers: Sequence[int], |
| name="pain", |
| log=None, |
| ): |
| layers = list(layers) |
| blocks = decoder_layers(model) |
| if ( |
| not layers |
| or len(set(layers)) != len(layers) |
| or any(i < 0 or i >= len(blocks) for i in layers) |
| ): |
| raise ValueError("Invalid extraction layer selection") |
| sums = {} |
| squares = {} |
| norms = {i: 0.0 for i in layers} |
| count = 0 |
| digest = hashlib.sha256() |
| was_training = model.training |
| model.eval() |
| try: |
| for pair in pairs: |
| if pair.get("split", "train") != "train": |
| raise ValueError("Extraction only accepts the training split") |
| target = pair.get("target", pair.get(name)) |
| if target is None: |
| raise ValueError(f"Missing contrastive answer for {name}") |
| if pair.get("channel", name) != name: |
| raise ValueError("Mixed channels in one extraction run") |
| record = { |
| "situation": pair["situation"], |
| "calm": pair["calm"], |
| "target": target, |
| } |
| digest.update( |
| (json.dumps(record, sort_keys=True, ensure_ascii=False) + "\n").encode() |
| ) |
| baseline = answer_statistics( |
| model, tokenizer, record["situation"], record["calm"], layers |
| ) |
| affected = answer_statistics( |
| model, tokenizer, record["situation"], target, layers |
| ) |
| count += 1 |
| for layer in layers: |
| delta = affected[layer][0] - baseline[layer][0] |
| sums[layer] = sums.get(layer, torch.zeros_like(delta)) + delta |
| squares[layer] = ( |
| squares.get(layer, torch.zeros_like(delta)) + delta.square() |
| ) |
| norms[layer] += (baseline[layer][1] + affected[layer][1]) * 0.5 |
| if log and count % 10 == 0: |
| log(f"{name}: {count} contrastive pairs") |
| finally: |
| model.train(was_training) |
| if not count: |
| raise ValueError("No training pairs") |
| directions = {} |
| stats = {} |
| for layer in layers: |
| mean = sums[layer] / count |
| magnitude = mean.norm().item() |
| if magnitude < 1e-8: |
| raise ValueError(f"No contrastive signal at layer {layer}") |
| directions[layer] = mean / magnitude |
| norms[layer] /= count |
| variance = (squares[layer] / count - mean.square()).clamp_min(0) |
| stats[str(layer)] = { |
| "mean_difference_norm": magnitude, |
| "coordinate_variance_mean": variance.mean().item(), |
| "standard_error_l2": (variance.sum() / count).sqrt().item(), |
| } |
| config = model.config |
| metadata = { |
| "release": "CODELOCK-8Bbeta-v0.2", |
| "signal": name, |
| "pairs": count, |
| "split": "train", |
| "base_model": getattr(config, "_name_or_path", "unknown"), |
| "model_revision": getattr(config, "_commit_hash", None), |
| "hidden_size": next(iter(directions.values())).numel(), |
| "num_hidden_layers": len(blocks), |
| "layers": layers, |
| "data_sha256": digest.hexdigest(), |
| "capture_site": "decoder block output, before any final model norm", |
| "method": "paired mean difference over answer tokens, unit direction per layer", |
| "reference_norm": "mean per-answer-token L2 activation norm", |
| "statistics": stats, |
| } |
| return AffectVector(directions, norms, metadata).validate() |
|
|
