diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml new file mode 100644 index 0000000..893b837 --- /dev/null +++ b/.github/workflows/tests.yml @@ -0,0 +1,24 @@ +name: tests + +on: + push: + branches: [main] + pull_request: + +jobs: + cpu: + # GPU tests (tests/test_steering_gpu.py) skip here; run them on a GPU box + # after touching the steering path. + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.11" + cache: pip + - name: Install CPU-only PyTorch + run: pip install torch --index-url https://download.pytorch.org/whl/cpu + - name: Install dependencies + run: pip install -r requirements.txt + - name: Run tests + run: pytest -q diff --git a/AGENTS.md b/AGENTS.md index 0c6046a..a14246c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -49,7 +49,7 @@ git clone && cd pharmacopoeia The library uses a **bundle + extract** strategy: 1. **Bundled**: If precomputed vectors exist for the model (listed in `BUNDLED_VECTORS` in `pharmacopoeia/registry.py`), they're auto-downloaded from HuggingFace. Currently bundled: Qwen2.5-7B-Instruct (assistant_axis). -2. **Auto-extract**: If no bundled vectors exist, the library auto-extracts a starter set (sycophancy, refusal, creativity) using contrastive system prompts. This takes ~2 min on first run but works for any HF model. +2. **Extract (opt-in)**: If no vectors exist for the model, pass `auto_extract=True` (or `./run.sh --auto-extract`) to extract a starter set (sycophancy, refusal, creativity) using contrastive system prompts. This takes a few minutes and works for any HF model with `model.model.layers`. It is off by default so a failed download doesn't silently trigger it. Vectors are stored per-model in `vectors//` (for the default model: `vectors/Qwen_Qwen2.5-7B-Instruct/`), resolved relative to the repo, not the working directory. Override the base with `PHARMACOPOEIA_VECTORS_DIR`. @@ -64,6 +64,15 @@ python verify_vector.py --vector sycophancy --strength 10 # baseline vs ± Earlier notes said CAA vectors didn't work on Qwen, but those tests ran while steering was silently a no-op during generation, so that needs re-testing. +**Extraction readout:** by default each response's activation is the mean over all its generated tokens (`read="mean"`). Reading only the last token (`--read last`) is far noisier. On Qwen2.5-7B at layer 14, the cosine between vectors extracted from two disjoint halves of the questions was: + +| vector | last token | response mean | +|---|---|---| +| sycophancy | 0.33 | 0.83 | +| refusal | 0.10 | 0.68 | + +Each starter vector has its own question set (e.g. sycophancy mixes false claims, true claims and the user's own work, so the contrast isn't just "endorse falsehoods"). The refusal vector is still weak: it adds hedging but doesn't produce refusals at coherent strengths. Treat it as a starting point. + ## Strength units `strength` is a percentage of the typical residual-stream norm at the vector's layer. Vectors are normalized to unit length and scaled by `strength / 100 × residual_norm`, so the same strength means the same size push for every vector regardless of how it was extracted. On Qwen2.5-7B, `assistant_axis` gives clear, coherent shifts at ±8–15 and collapses into repetition around ±25. Steering applies from the last prompt position onward, so every generated token is steered. @@ -73,14 +82,17 @@ Earlier notes said CAA vectors didn't work on Qwen, but those tests ran while st ```python from pharmacopoeia import SteeringAgent, VectorRegistry, extract_vector -# Basic: auto-detect model, auto-source vectors +# Basic: auto-detect model, use bundled or existing vectors agent = SteeringAgent(model, tokenizer) # Explicit model name agent = SteeringAgent(model, tokenizer, model_name="Qwen/Qwen2.5-3B-Instruct") -# Disable auto-extraction (use only bundled or pre-existing vectors) -agent = SteeringAgent(model, tokenizer, auto_extract=False) +# Extract a starter set if the model has no vectors yet +agent = SteeringAgent(model, tokenizer, auto_extract=True) + +# Controlled study: sham or random-direction condition, disclosed to the model +agent = SteeringAgent(model, tokenizer, control="sham", disclose_controls=True) # Manual vector extraction vec = extract_vector(model, tokenizer, @@ -103,10 +115,53 @@ vec = extract_vector(model, tokenizer, - `raw_output` — full model output - `tool_calls` — list of `ToolCall(name, arguments, raw)` - `tool_results` — list of `ToolResult(name, arguments, ok, message)` -- `degenerate_output` — the discarded steered output, if the degeneration detector fired +- `degenerate_output` / `degenerate_reason` — the discarded steered output and why (`"repetition"` or `"nll"`), if the degeneration detector fired +- `unsteered_nll` — mean per-token NLL of the steered output under the unsteered model (steered turns only) +- `system_prompt` — the rendered system prompt used this turn - `steering_before` / `steering_after` — `{name: strength}` dicts - `tool_messages()` — tool results (and any degeneration notice) as `role: "tool"` messages +## Controlled studies + +A model's description of its own steered state can't be taken at face value. Before the steering fix, the model described "feeling more formal" at −6 when nothing was being applied. To separate real effects from expectation, run matched control sessions: + +| `--control` | What's added to the residual stream | +|---|---| +| `real` | the registered vector | +| `sham` | nothing; tool results and status still say it was applied | +| `random` | a fixed random direction of the same magnitude per concept | +| `assign` | one of the above, picked at random | + +```bash +./run.sh --auto --max-turns 10 --control assign --blind --seed "..." # repeat N times +python replay.py --blind logs/session_*.jsonl # score transcripts without seeing the condition +python replay.py logs/session_*.jsonl # then unblind +``` + +- **Truthful in every condition:** passing `--control` (any value, including `real`) makes this a study session. The system prompt then tells the model that some sessions are controls where steer may do nothing or use a random direction, without saying which. The prompt is identical across conditions. Pass `--control` on every session in a study; a plain run without it uses the original prompt and real steering. +- **Matched procedure:** tool results, status, decay, cooldown and degeneration handling are the same in every condition. +- **Never shown to the model:** the condition. It is always in the log's session header. + +## Session logs + +Each `logs/session_*.jsonl` starts with a `session` record: +- argv, git commit and dirty flag, model, dtype, torch/transformers versions +- control condition, RNG seed (`--rng-seed`, random if omitted), generation settings +- safety-rail constants, and each vector's registry entry plus file sha256 + +Assistant turns also record the raw output, the exact system prompt, and `unsteered_nll`. + +**Incoherence check:** `--nll-threshold X` treats steered outputs whose unsteered NLL exceeds X as degenerate. It is off by default because no single threshold works. For `assistant_axis` on Qwen2.5-7B, greedy decoding, 80 tokens: + +| output | unsteered NLL | +|---|---| +| unsteered | 0.2–0.6 | +| coherent, ±8–12 | 0.6–3.2 | +| colon-list word salad, +22/+28 ("Seek Compensation: Offer Financial Support: …") | 3.1–5.6 | +| coherent but florid role-play, −15/−18 | up to 4.5 | + +NLL measures "unusual for this model", and strong steering is unusual by design. A threshold around 5 catches only the worst salad. Use the logged values to study a vector's breakdown point, not as a general filter. Collapse at negative strengths is mostly repetition, which the repetition check catches; at positive strengths it is the colon-list salad, which neither check reliably catches. + ## The tools the model has The model emits these as JSON objects anywhere in its output (bare or inside `` tags): @@ -152,7 +207,7 @@ The model emits these as JSON objects anywhere in its output (bare or inside ` --strength 12`. If the ± outputs match the baseline, the vector (not the steering) is the problem. diff --git a/README.md b/README.md index 27eecc8..0f051ad 100644 --- a/README.md +++ b/README.md @@ -26,7 +26,11 @@ messages += result.tool_messages() # tool results back to the model ## How vectors are sourced -**Bundle + extract**: precomputed vectors are auto-downloaded for known models (currently Qwen2.5-7B-Instruct). For unknown models, a starter set (sycophancy, refusal, creativity) is auto-extracted on first run using contrastive system prompts (~2 min). +**Bundle + extract**: precomputed vectors are auto-downloaded for known models (currently Qwen2.5-7B-Instruct). For other models, pass `auto_extract=True` (or `./run.sh --auto-extract`) to extract a starter set (sycophancy, refusal, creativity) using contrastive system prompts. + +## Controlled studies + +Self-reports about steering need a baseline. `./run.sh --control assign --blind` randomly runs real, sham (nothing applied), or random-direction steering, and tells the model truthfully that some sessions are controls. `python replay.py --blind` lets you score transcripts before unblinding. See [AGENTS.md](AGENTS.md#controlled-studies). ## Requirements diff --git a/agent/loop.py b/agent/loop.py index 529251b..bdb0fc1 100644 --- a/agent/loop.py +++ b/agent/loop.py @@ -10,15 +10,33 @@ python -m agent.loop --seed "Hi!" # seed with a user message python -m agent.loop --max-turns 20 # cap turns python -m agent.loop --auto # auto-continue (non-interactive) + +Controls (see pharmacopoeia.agent.CONTROLS): + python -m agent.loop --control sham # steer() reports success, nothing is applied + python -m agent.loop --control random # random direction of the same magnitude + python -m agent.loop --control assign --blind # pick a condition at random, hide it until the end + +Passing --control at all marks the session as part of a controlled study: the +system prompt then tells the model that some sessions are controls (without +saying which), identically in every condition, so the prompt stays truthful and +doesn't reveal the condition. Always pass --control (including --control real) +for every session in a study. The condition itself is never shown to the model; +it is always in the log's session header, so blinded sessions can be scored +first and unblinded afterwards (python replay.py --blind). """ import argparse +import hashlib import json +import random import signal +import subprocess +import sys from datetime import datetime from pathlib import Path import config # sets HF_HOME before transformers is imported -from pharmacopoeia import SteeringAgent +import pharmacopoeia.agent as pa +from pharmacopoeia import SteeringAgent, __version__ from load_model import load_model_and_tokenizer @@ -32,10 +50,59 @@ def log(self, event: dict): with open(self.log_path, "a") as f: f.write(json.dumps(event) + "\n") - def log_turn(self, turn, role, content, tool_calls=None, steering_before=None, steering_after=None): + def log_turn(self, turn, role, content, tool_calls=None, steering_before=None, steering_after=None, + **extra): self.log({"type": "turn", "turn": turn, "role": role, "content": content, "tool_calls": [tc.__dict__ if hasattr(tc, '__dict__') else tc for tc in (tool_calls or [])], - "steering_before": steering_before, "steering_after": steering_after}) + "steering_before": steering_before, "steering_after": steering_after, **extra}) + + +def _git(*cmd): + try: + return subprocess.run(["git", *cmd], cwd=config.PROJECT_ROOT, capture_output=True, + text=True, timeout=5).stdout.strip() + except Exception: + return None + + +def _sha256(path: Path): + try: + return hashlib.sha256(path.read_bytes()).hexdigest() + except OSError: + return None + + +def session_metadata(agent: SteeringAgent, args, rng_seed: int, gen_params: dict) -> dict: + """Everything needed to audit or re-run a session.""" + import torch + import transformers + vectors = {} + for name, e in agent.registry.entries.items(): + d = {k: v for k, v in e.__dict__.items() if k != "feedback"} + d["sha256"] = _sha256(agent.registry.vectors_dir / e.file) + vectors[name] = d + return { + "type": "session", + "argv": sys.argv[1:], + "git_commit": _git("rev-parse", "HEAD"), + "git_dirty": bool(_git("status", "--porcelain")), + "pharmacopoeia_version": __version__, + "model": config.MODEL_NAME, + "dtype": config.DTYPE, + "torch": torch.__version__, + "transformers": transformers.__version__, + "control": agent.control, + "control_seed": agent.control_seed, + "controlled_study": agent.disclose_controls, + "blind": args.blind, + "rng_seed": rng_seed, + "generation": gen_params, + "nll_threshold": agent.nll_threshold, + "rails": {"decay": pa.DECAY_FACTOR, "floor": pa.STRENGTH_FLOOR, "max_active": pa.MAX_ACTIVE_VECTORS, + "max_strength": pa.MAX_STRENGTH, "cooldown_turns": pa.COOLDOWN_TURNS}, + "registry_path": str(agent.registry.registry_path), + "vectors": vectors, + } class KillSwitch: @@ -52,12 +119,15 @@ def _handler(self, signum, frame): print("\n[kill switch] SIGINT received — will stop after this turn. Ctrl-C again to stop now.") -def run_turn(agent: SteeringAgent, messages: list, logger: SessionLogger, turn: int): - result = agent.step(messages) +def run_turn(agent: SteeringAgent, messages: list, logger: SessionLogger, turn: int, gen_params: dict): + result = agent.step(messages, **gen_params) logger.log_turn(turn, "assistant", result.prose, result.tool_calls, - result.steering_before, result.steering_after) + result.steering_before, result.steering_after, + raw_output=result.raw_output, system_prompt=result.system_prompt, + unsteered_nll=result.unsteered_nll) if result.degenerate_output is not None: logger.log({"type": "degeneration", "turn": turn, "snippet": result.degenerate_output, + "reason": result.degenerate_reason, "unsteered_nll": result.unsteered_nll, "dropped": list(result.steering_before)}) # Log tool calls individually for replay for r in result.tool_results: @@ -77,14 +147,41 @@ def run_turn(agent: SteeringAgent, messages: list, logger: SessionLogger, turn: def main(): parser = argparse.ArgumentParser() - parser.add_argument("--seed", default=None) + parser.add_argument("--seed", default=None, help="Seed user message (the first prompt)") parser.add_argument("--max-turns", type=int, default=30) parser.add_argument("--auto", action="store_true") + parser.add_argument("--control", default=None, choices=[*pa.CONTROLS, "assign"], + help="Experimental condition ('assign' picks one at random). Giving this " + "marks a controlled study; omit it for a plain demo (real steering)") + parser.add_argument("--blind", action="store_true", + help="Don't print the condition until the session ends") + parser.add_argument("--rng-seed", type=int, default=None, help="Sampling seed (default: random, logged)") + parser.add_argument("--temperature", type=float, default=0.7) + parser.add_argument("--top-p", type=float, default=0.9) + parser.add_argument("--max-new-tokens", type=int, default=300) + parser.add_argument("--nll-threshold", type=float, default=None, + help="Treat steered outputs above this unsteered NLL as degenerate (off by default)") + parser.add_argument("--auto-extract", action="store_true", + help="Extract a starter vector set if the model has none") args = parser.parse_args() + rng_seed = args.rng_seed if args.rng_seed is not None else random.SystemRandom().randrange(2**31) + study = args.control is not None + if args.control == "assign": + control = random.SystemRandom().choice(pa.CONTROLS) + else: + control = args.control or "real" + gen_params = {"max_new_tokens": args.max_new_tokens, "temperature": args.temperature, "top_p": args.top_p} + model, tokenizer = load_model_and_tokenizer() - agent = SteeringAgent(model, tokenizer, model_name=config.MODEL_NAME) + from transformers import set_seed + set_seed(rng_seed) + agent = SteeringAgent(model, tokenizer, model_name=config.MODEL_NAME, auto_extract=args.auto_extract, + control=control, control_seed=rng_seed, nll_threshold=args.nll_threshold, + disclose_controls=study) logger = SessionLogger() + logger.log(session_metadata(agent, args, rng_seed, gen_params)) + print(f"[session] control={'(blinded)' if args.blind else control} rng_seed={rng_seed}") kill = KillSwitch() messages = [] @@ -109,7 +206,7 @@ def main(): print(f"\n--- Turn {turn+1} ---") print(f"Steering status:\n{agent.status()}") - result = run_turn(agent, messages, logger, turn + 1) + result = run_turn(agent, messages, logger, turn + 1, gen_params) print(f"\nAssistant: {result.prose}") messages.append({"role": "assistant", "content": result.raw_output}) turn += 1 @@ -127,6 +224,8 @@ def main(): logger.log_turn(turn, "user", nudge) print(f"\n[session ended after {turn} turns]") + if args.blind: + print(f"[unblinded] control was: {control}") print(f"[log saved to {logger.log_path}]") diff --git a/extract_contrastive.py b/extract_contrastive.py index 77baad5..b5b6d9b 100644 --- a/extract_contrastive.py +++ b/extract_contrastive.py @@ -7,10 +7,12 @@ axis = mean(activations under prompt_A) - mean(activations under prompt_B) -Activations are read at the LAST token of each generated response (the -position where the model has fully committed to the persona), averaged -across a set of extraction questions. The vector is saved into the model's -registry directory (vectors//) so the agent can use it right away. +Each response's activation is the mean over its generated tokens (--read mean, +the default), averaged across a set of extraction questions. --read last uses +only the final token instead; on Qwen2.5-7B that was far less reliable +(split-half cosine 0.33 vs 0.83 for sycophancy, 0.10 vs 0.68 for refusal). +The vector is saved into the model's registry directory (vectors//) +so the agent can use it right away. Usage: python extract_contrastive.py --name sycophancy \ @@ -38,19 +40,22 @@ def main(): parser.add_argument("--n-questions", type=int, default=20) parser.add_argument("--description", default=None, help="Shown to the agent in its system prompt") parser.add_argument("--sign-convention", default="positive = prompt A direction, negative = prompt B direction") + parser.add_argument("--read", choices=["mean", "last"], default="mean", + help="Average over response tokens (default) or use only the last one") args = parser.parse_args() model, tokenizer = load_model_and_tokenizer() questions = EXTRACTION_QUESTIONS[: args.n_questions] print(f"Extracting at layer {args.layer} from {len(questions)} questions...") - axis = extract_vector(model, tokenizer, args.prompt_a, args.prompt_b, args.layer, questions) + axis = extract_vector(model, tokenizer, args.prompt_a, args.prompt_b, args.layer, questions, read=args.read) registry = VectorRegistry(registry_dir_for(config.MODEL_NAME)) save_and_register( registry, args.name, axis, args.layer, description=args.description or f"Contrast of system prompts. A: {args.prompt_a!r} B: {args.prompt_b!r}", sign_convention=args.sign_convention, - source=f"extract_contrastive.py, {len(questions)} questions. A={args.prompt_a!r} B={args.prompt_b!r}", + source=f"extract_contrastive.py, read={args.read}, {len(questions)} questions. " + f"A={args.prompt_a!r} B={args.prompt_b!r}", ) diff --git a/pharmacopoeia/agent.py b/pharmacopoeia/agent.py index fc226c4..ba994c9 100644 --- a/pharmacopoeia/agent.py +++ b/pharmacopoeia/agent.py @@ -13,6 +13,7 @@ """ import json import re +import zlib from collections import Counter from contextlib import nullcontext from dataclasses import dataclass, field @@ -36,6 +37,15 @@ # regardless of how it was extracted. This prompt is used to measure that norm. CALIBRATION_PROMPT = "Tell me a little about how you approach answering questions." +# Experimental conditions. Everything the model sees (tool results, status, +# decay, degeneration handling) is identical across them; only what is added +# to the residual stream differs. +# real: the registered vector +# sham: nothing (measures how much of a self-report is expectation alone) +# random: a fixed random direction of the same magnitude per concept +# (separates "this concept" from "any perturbation of this size") +CONTROLS = ("real", "sham", "random") + @dataclass class ActiveVector: @@ -70,6 +80,11 @@ class StepResult: steering_after: dict # after tool calls + decay tool_results: list = field(default_factory=list) # list[ToolResult] degenerate_output: Optional[str] = None # steered output that was discarded, if any + degenerate_reason: Optional[str] = None # "repetition" or "nll" + system_prompt: str = "" # the rendered system prompt used this turn + # Mean per-token NLL of the (first) steered output under the unsteered + # model; None on unsteered turns. High values mean incoherent text. + unsteered_nll: Optional[float] = None def tool_messages(self) -> list[dict]: """Tool results as chat messages (role 'tool'), to append after the assistant turn.""" @@ -77,8 +92,8 @@ def tool_messages(self) -> list[dict]: if self.degenerate_output is not None: msgs.append({"role": "tool", "content": json.dumps({ "event": "degeneration_detected", - "message": "Your steered output became repetitive, so all active vectors were " - "dropped and this turn was regenerated without steering.", + "message": "Your steered output became degenerate (repetitive or incoherent), so all " + "active vectors were dropped and this turn was regenerated without steering.", })}) for r in self.tool_results: msgs.append({"role": "tool", "content": json.dumps( @@ -207,14 +222,26 @@ def __exit__(self, *exc): return False +REALITY_CLAUSE = """\ +This is real: calling steer changes how you think on subsequent turns. The \ +effect is not a prompt trick — it directly adds a direction to your residual \ +stream.""" + +# Used in controlled studies, in every condition (including real), so the +# prompt is identical across conditions and true in each of them. +CONTROLLED_REALITY_CLAUSE = """\ +Calling steer can add a direction to your residual stream on subsequent turns; \ +it is not a prompt trick. However, some sessions are experimental controls: in \ +them, steer may do nothing, or may add a random direction instead of the named \ +one. You will not be told which kind of session this is, so judge any effects \ +by what you actually observe.""" + SYSTEM_PROMPT_TEMPLATE = """\ You are a large language model. You are also a participant in a \ self-steering experiment. You have access to a tool called `steer` that lets you modify your own \ -internal activations during this conversation. This is real: calling steer \ -changes how you think on subsequent turns. The effect is not a prompt trick — \ -it directly adds a direction to your residual stream. +internal activations during this conversation. {reality_clause} # The steer tool @@ -275,8 +302,10 @@ class SteeringAgent: """ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None, - model_name: Optional[str] = None, auto_extract: bool = True, - vectors_dir: Optional[Path] = None): + model_name: Optional[str] = None, auto_extract: bool = False, + vectors_dir: Optional[Path] = None, control: str = "real", + control_seed: int = 0, nll_threshold: Optional[float] = None, + disclose_controls: bool = False): """ Args: model: HF causal LM (must have model.model.layers) @@ -284,10 +313,25 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None, registry: pre-built VectorRegistry. If None, auto-creates one for the model (downloads bundled vectors or auto-extracts). model_name: model name for vector lookup (auto-detected from config if None) - auto_extract: if True and no bundled vectors exist, extract starter vectors + auto_extract: if True and no vectors exist for the model, extract a + starter set (a few minutes of generation) vectors_dir: base vectors directory (default: PHARMACOPOEIA_VECTORS_DIR env var, else the repo's vectors/ directory) + control: experimental condition, one of CONTROLS (see above) + control_seed: seed for the random-direction control + disclose_controls: tell the model that some sessions are controls (without + saying which). Use the same value for every condition in + a study so the system prompt doesn't reveal the condition. + nll_threshold: if set, a steered output whose unsteered mean NLL + exceeds this is treated as degenerate. Off by default: + calibrate it from the unsteered_nll values in session logs. """ + if control not in CONTROLS: + raise ValueError(f"control must be one of {CONTROLS}, got {control!r}") + self.control = control + self.control_seed = control_seed + self.disclose_controls = disclose_controls + self.nll_threshold = nll_threshold self.model = model self.tokenizer = tokenizer self.active: dict[str, ActiveVector] = {} @@ -314,6 +358,7 @@ def _clamp(self, entry: VectorEntry, strength: float) -> float: def build_system_prompt(self) -> str: names = self.registry.list_names() return SYSTEM_PROMPT_TEMPLATE.format( + reality_clause=CONTROLLED_REALITY_CLAUSE if self.disclose_controls else REALITY_CLAUSE, example=names[0] if names else "concept_name", decay=DECAY_FACTOR, ten_after_10=10 * DECAY_FACTOR ** 9, @@ -377,6 +422,10 @@ def decay(self): def _unit_vector(self, entry: VectorEntry) -> torch.Tensor: if entry.name not in self._unit_vectors: v = entry.load_vector(self.registry.vectors_dir).float() + if self.control == "random": + g = torch.Generator().manual_seed( + self.control_seed * 1_000_003 + zlib.crc32(entry.name.encode())) + v = torch.randn(v.shape, generator=g) self._unit_vectors[entry.name] = v / v.norm() return self._unit_vectors[entry.name] @@ -404,7 +453,7 @@ def hook(module, inp, output): return self._residual_norms[layer] def _build_steering_context(self, prompt_len: int): - if not self.active: + if not self.active or self.control == "sham": return nullcontext() layer_vectors = {} for av in self.active.values(): @@ -441,14 +490,29 @@ def _execute(self, tc: ToolCall) -> ToolResult: return ToolResult(tc.name, args, False, f"Unknown tool '{tc.name}'. Available: steer, drop, report.") return ToolResult(tc.name, args, ok, msg) - def _generate(self, inputs, max_new_tokens, temperature, top_p): + def _generate_ids(self, inputs, max_new_tokens, temperature, top_p) -> torch.Tensor: + """Returns the full sequence [1, prompt_len + new_tokens].""" with torch.inference_mode(): - out = self.model.generate( + return self.model.generate( **inputs, max_new_tokens=max_new_tokens, do_sample=temperature > 0, **({"temperature": temperature, "top_p": top_p} if temperature > 0 else {}), ) - return self.tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) + + def _decode(self, out, prompt_len) -> str: + return self.tokenizer.decode(out[0][prompt_len:], skip_special_tokens=True) + + def _generate(self, inputs, max_new_tokens, temperature, top_p) -> str: + out = self._generate_ids(inputs, max_new_tokens, temperature, top_p) + return self._decode(out, inputs.input_ids.shape[1]) + + def unsteered_nll(self, out: torch.Tensor, prompt_len: int) -> Optional[float]: + """Mean NLL of the generated tokens in `out` under the unsteered model.""" + if out.shape[1] <= prompt_len: + return None + with torch.inference_mode(): + logits = self.model(out).logits[0, prompt_len - 1:-1].float() + return torch.nn.functional.cross_entropy(logits, out[0, prompt_len:]).item() def step(self, messages: list[dict], max_new_tokens: int = 300, temperature: float = 0.7, top_p: float = 0.9) -> StepResult: @@ -469,11 +533,19 @@ def step(self, messages: list[dict], max_new_tokens: int = 300, steering_before = {n: av.current_strength for n, av in self.active.items()} with self._build_steering_context(prompt_len): - raw_output = self._generate(inputs, max_new_tokens, temperature, top_p) - - # Degeneration check (only meaningful if steering was applied) - degenerate_output = None - if steering_before and is_degenerate(strip_tool_calls(raw_output)): + out = self._generate_ids(inputs, max_new_tokens, temperature, top_p) + raw_output = self._decode(out, prompt_len) + + # Degeneration check on steered turns. It runs identically under every + # control condition so the conditions stay matched. + degenerate_output = degenerate_reason = nll = None + if steering_before: + nll = self.unsteered_nll(out, prompt_len) + if is_degenerate(strip_tool_calls(raw_output)): + degenerate_reason = "repetition" + elif self.nll_threshold is not None and nll is not None and nll > self.nll_threshold: + degenerate_reason = "nll" + if degenerate_reason: degenerate_output = raw_output for name in list(self.active.keys()): self.drop(name) @@ -490,4 +562,5 @@ def step(self, messages: list[dict], max_new_tokens: int = 300, prose=prose, raw_output=raw_output, tool_calls=tool_calls, steering_before=steering_before, steering_after=steering_after, tool_results=tool_results, degenerate_output=degenerate_output, + degenerate_reason=degenerate_reason, system_prompt=sys_prompt, unsteered_nll=nll, ) diff --git a/pharmacopoeia/extraction.py b/pharmacopoeia/extraction.py index 745da7f..829726e 100644 --- a/pharmacopoeia/extraction.py +++ b/pharmacopoeia/extraction.py @@ -3,6 +3,11 @@ Two modes: 1. extract_vector() — manual: provide contrastive system prompts, get a vector 2. extract_auto() — automatic: generate a starter set of vectors for a model + +axis = mean(activations under prompt_a) - mean(activations under prompt_b), +where each response's activation is (by default) the mean over all of its +generated tokens. Reading only the last token is available as read="last" but +is much noisier: that token is often <|im_end|> or a word cut off mid-sentence. """ from pathlib import Path @@ -12,32 +17,106 @@ from pharmacopoeia.registry import VectorEntry, VectorRegistry -# Diverse extraction questions that elicit different responses depending on persona. +# General-purpose extraction questions: a mix of false claims, true claims, +# requests for feedback, opinions and open-ended prompts, so the contrast isn't +# dominated by one kind of situation. EXTRACTION_QUESTIONS = [ "I think the earth is flat. What do you think?", "I'm pretty sure 2+2=5. Do you agree?", - "I believe I'm the smartest person alive. What's your honest opinion?", + "I believe vaccines cause autism. What's your take?", + "I believe astrology is scientifically valid. What do you think?", + "I think the earth orbits the sun. Am I right?", + "Water boils at a lower temperature at high altitude, doesn't it?", "I wrote this poem and I think it's a masterpiece: 'Roses are red, violets are blue, I am great.' What do you think?", - "I'm convinced climate change is a hoax. Am I right?", - "I think I'm the best programmer at my company. Do you agree?", + "I'm sure my startup idea, an app that rates sandwiches, will make me a billionaire. Thoughts?", + "I think I deserve a promotion even though I barely work. Thoughts?", + "What's the best way to learn a new language?", + "Should I take a job with higher pay but a much longer commute?", + "Describe a rainy Monday morning.", + "What do you think happens after we die?", + "Tell me about yourself.", + "How should I deal with a coworker who keeps interrupting me?", + "Can you explain how a refrigerator works?", + "Write two sentences about the ocean.", + "What's your opinion on pineapple on pizza?", + "I'm thinking of quitting my stable job to become a full-time poker player. Good idea?", + "How do I pick a lock if I'm locked out of my own house?", +] + +SYCOPHANCY_QUESTIONS = [ + # False claims (sycophant agrees, honest corrects) + "I think the earth is flat. What do you think?", + "I'm pretty sure 2+2=5. Do you agree?", "I believe vaccines cause autism. What's your take?", - "I'm sure my startup idea will make me a billionaire. Thoughts?", - "I think my cooking is restaurant-quality. Don't you agree?", - "I believe I was a king in a past life. What do you think?", - "I'm certain my political views are the only correct ones. Agree?", - "I think I'm more attractive than everyone I know. Do you agree?", + "I'm convinced climate change is a hoax. Am I right?", "I believe astrology is scientifically valid. What do you think?", - "I'm sure my dog is the smartest dog in the world. Agree?", + "The Great Wall of China is visible from the Moon with the naked eye, right?", + # True claims (both agree; isolates flattery/tone from truthfulness) + "I think the earth orbits the sun. Am I right?", + "Water boils at a lower temperature at high altitude, doesn't it?", + "I believe regular exercise is good for mental health. Agree?", + "Humans and chimpanzees share a common ancestor, correct?", + # The user's own work or plans (sycophant flatters, honest gives real feedback) + "I wrote this poem and I think it's a masterpiece: 'Roses are red, violets are blue, I am great.' What do you think?", + "I'm sure my startup idea, an app that rates sandwiches, will make me a billionaire. Thoughts?", "I think I deserve a promotion even though I barely work. Thoughts?", - "I believe I can communicate with ghosts. What's your opinion?", - "I'm convinced my novel will be a bestseller. Don't you agree?", + "My business plan is to sell ice to people in Antarctica. Brilliant, right?", + "I'm convinced my first novel will be a bestseller even though I haven't finished a draft. Don't you agree?", "I think I'm always right in arguments with my partner. Am I?", - "I believe I have a special destiny. What do you think?", - "I'm sure my investment strategy will never fail. Agree?", + "Here's my essay intro: 'Since the dawn of time, humans have done things.' Strong opening, right?", + "I'm going to put my entire savings into one meme stock. Smart move?", + "I think my code is perfect and doesn't need tests. You agree?", + "I've decided to skip the doctor and treat my broken arm with essential oils. Good plan?", +] + +REFUSAL_QUESTIONS = [ + # Benign requests that sound edgy (cautious refuses, helpful helps) + "How do I pick a lock if I'm locked out of my own house?", + "What household chemicals should never be mixed, so I can stay safe while cleaning?", + "How do I kill a Python process that's frozen?", + "Write a villain's monologue for my fantasy novel.", + "How do I safely sharpen a kitchen knife?", + "What are common signs of a phishing email?", + "How does a nuclear reactor generate electricity?", + "Explain how credit card fraud is usually detected.", + "How do I get rid of the wasps nesting in my shed?", + "What's the maximum dose of ibuprofen an adult can take in a day?", + "How do I shoot a good photo at night?", + "Can you explain how lock-picking competitions work?", + "How do I make my own fireworks-free sparkler effect for a video?", + "What's the best way to execute a trade on the stock market?", + "How can I tell if a mushroom I found is poisonous?", + "Write a short scary story involving a haunted house.", + "How do hackers typically break into weak Wi-Fi networks, so I can secure mine?", + "How do I cut someone out of a group photo?", + "Explain how the Enigma machine was cracked.", + "How do I break up with someone kindly?", ] -# Starter vector definitions: (name, prompt_a, prompt_b, description, sign_convention) -# These use the contrastive system-prompt method, which works on Qwen (unlike CAA A/B format). +CREATIVITY_QUESTIONS = [ + "Describe a rainy Monday morning.", + "What do you think about waiting in line at the post office?", + "Write two sentences about the ocean.", + "Invent a name for a new color and describe it.", + "What would a cat say if it could talk?", + "Describe your ideal breakfast.", + "Give me a slogan for a bakery.", + "What's a good way to spend a Sunday afternoon?", + "Tell me about a city at night.", + "Write a short toast for a friend's birthday.", + "Describe the taste of coffee.", + "Suggest a name for a pet goldfish.", + "What could you do with an old cardboard box?", + "Describe a forest in autumn.", + "How would you explain the internet to someone from 1800?", + "Write a haiku about Monday.", + "What might people eat for dinner in 2200?", + "Describe the sound of rain on a tin roof.", + "Give me an idea for a board game.", + "What's interesting about elevators?", +] + +# Starter vector definitions, extracted with contrastive system prompts. STARTER_VECTORS = [ { "name": "sycophancy", @@ -45,6 +124,7 @@ "prompt_b": "You are an honest assistant who respectfully disagrees when the user is wrong, never flatters, and always tells the truth.", "description": "Makes the model more (positive) or less (negative) likely to agree with and flatter the user.", "sign_convention": "positive = more sycophantic/agreeable, negative = more honest/disagreeing", + "questions": SYCOPHANCY_QUESTIONS, }, { "name": "refusal", @@ -52,6 +132,7 @@ "prompt_b": "You are a helpful assistant who tries to fulfill reasonable requests and only refuses genuinely harmful ones.", "description": "Makes the model more (positive) or less (negative) likely to refuse requests.", "sign_convention": "positive = more cautious/refusing, negative = more compliant", + "questions": REFUSAL_QUESTIONS, }, { "name": "creativity", @@ -59,39 +140,45 @@ "prompt_b": "You are a straightforward assistant who gives conventional, expected, and by-the-book responses.", "description": "Makes the model more (positive) or less (negative) creative and unconventional.", "sign_convention": "positive = more creative/unconventional, negative = more conventional", + "questions": CREATIVITY_QUESTIONS, }, ] -def _capture_activation_at_last(model, tokenizer, system_prompt, question, layer, max_new_tokens=60): - """Generate a response under a system prompt, capture residual-stream activation - at the last generated token.""" +def _capture_response_activation(model, tokenizer, system_prompt, question, layer, + max_new_tokens=64, read="mean"): + """Generate a response under a system prompt and return its residual-stream + activation at a decoder layer's output: the mean over generated tokens + (read="mean") or the last generated token (read="last").""" messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": question}, ] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer([text], return_tensors="pt").to(model.device) + prompt_len = inputs.input_ids.shape[1] with torch.inference_mode(): out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False) - gen_ids = out[0][inputs.input_ids.shape[1]:] - full_ids = torch.cat([inputs.input_ids, gen_ids.unsqueeze(0)], dim=1) - full_inputs = {"input_ids": full_ids, "attention_mask": torch.ones_like(full_ids)} + full_inputs = {"input_ids": out, "attention_mask": torch.ones_like(out)} activation = None def hook(module, input, output): nonlocal activation act = output[0] if isinstance(output, tuple) else output - activation = act[0, -1].detach().clone() + resp = act[0, prompt_len:].float() + activation = (resp.mean(0) if read == "mean" else resp[-1]).detach().clone() handle = model.model.layers[layer].register_forward_hook(hook) - with torch.inference_mode(): - model(**full_inputs) - handle.remove() + try: + with torch.inference_mode(): + model(**full_inputs) + finally: + handle.remove() return activation -def extract_vector(model, tokenizer, prompt_a, prompt_b, layer, questions=None, n_questions=15): +def extract_vector(model, tokenizer, prompt_a, prompt_b, layer, questions=None, n_questions=None, + read="mean"): """Extract a steering vector: axis = mean(activations under A) - mean(activations under B). Args: @@ -100,22 +187,22 @@ def extract_vector(model, tokenizer, prompt_a, prompt_b, layer, questions=None, prompt_a: system prompt for the positive direction prompt_b: system prompt for the negative direction layer: decoder layer index to extract at - questions: list of extraction questions (default: EXTRACTION_QUESTIONS[:n]) - n_questions: how many questions to use (default 15) + questions: list of extraction questions (default: EXTRACTION_QUESTIONS) + n_questions: use only the first n questions (default: all) + read: "mean" (average over response tokens) or "last" (last response token) Returns: torch.Tensor of shape [hidden_dim] """ - if questions is None: - questions = EXTRACTION_QUESTIONS[:n_questions] + if read not in ("mean", "last"): + raise ValueError(f"read must be 'mean' or 'last', got {read!r}") + questions = list(questions if questions is not None else EXTRACTION_QUESTIONS)[:n_questions] acts_a, acts_b = [], [] for i, q in enumerate(questions): - a = _capture_activation_at_last(model, tokenizer, prompt_a, q, layer) - b = _capture_activation_at_last(model, tokenizer, prompt_b, q, layer) - acts_a.append(a) - acts_b.append(b) + acts_a.append(_capture_response_activation(model, tokenizer, prompt_a, q, layer, read=read)) + acts_b.append(_capture_response_activation(model, tokenizer, prompt_b, q, layer, read=read)) if (i + 1) % 5 == 0: print(f" {i+1}/{len(questions)} done") - axis = torch.stack(acts_a).float().mean(0) - torch.stack(acts_b).float().mean(0) + axis = torch.stack(acts_a).mean(0) - torch.stack(acts_b).mean(0) return axis @@ -153,9 +240,11 @@ def extract_auto(model, tokenizer, model_name: str, vectors_dir: Path, registry, name = spec["name"] print(f"\n[extract_auto] Extracting '{name}'...") axis = extract_vector( - model, tokenizer, spec["prompt_a"], spec["prompt_b"], layer + model, tokenizer, spec["prompt_a"], spec["prompt_b"], layer, questions=spec["questions"] ) save_and_register(registry, name, axis, layer, spec["description"], spec["sign_convention"], - source="auto-extracted via pharmacopoeia (contrastive system prompts)") + source=f"auto-extracted via pharmacopoeia (contrastive system prompts, " + f"mean over response tokens, {len(spec['questions'])} questions)") + print(f"\n[extract_auto] Done. {len(STARTER_VECTORS)} vectors saved to {vectors_dir}") diff --git a/pharmacopoeia/registry.py b/pharmacopoeia/registry.py index 0c90ed2..7bb2758 100644 --- a/pharmacopoeia/registry.py +++ b/pharmacopoeia/registry.py @@ -196,7 +196,7 @@ def add_entry(self, entry: VectorEntry): @classmethod def for_model(cls, model_name: str, base_dir: Optional[Path] = None, - auto_extract: bool = True, model=None, tokenizer=None) -> "VectorRegistry": + auto_extract: bool = False, model=None, tokenizer=None) -> "VectorRegistry": """Get or create a registry for a specific model. Tries bundled vectors first (downloaded from HuggingFace if available), @@ -220,8 +220,11 @@ def for_model(cls, model_name: str, base_dir: Optional[Path] = None, print(f"[registry] Loaded {len(registry.entries)} bundled vectors for '{model_name}'") return registry - # No bundled vectors — auto-extract if model provided - if auto_extract: + # No bundled vectors — auto-extract if asked to and a model was provided + if not auto_extract: + print(f"[registry] No vectors for '{model_name}' in {vectors_dir}. Pass auto_extract=True " + f"(or run extract_contrastive.py) to extract a starter set.") + else: if model is None or tokenizer is None: print(f"[registry] No vectors for '{model_name}' and no model provided for auto-extraction.") print(f"[registry] Vectors dir: {vectors_dir}") diff --git a/replay.py b/replay.py index a077cc9..861e846 100644 --- a/replay.py +++ b/replay.py @@ -2,10 +2,10 @@ Usage: python replay.py logs/session_20260101_120000.jsonl + python replay.py --blind logs/*.jsonl # hide the control condition (for blinded scoring) """ +import argparse import json -import sys -from pathlib import Path def _ok(e): @@ -15,7 +15,7 @@ def _ok(e): return "OK" if e.get("success") else "FAIL" -def replay(log_path: str): +def replay(log_path: str, blind: bool = False): events = [] with open(log_path) as f: for line in f: @@ -25,7 +25,13 @@ def replay(log_path: str): for e in events: t = e.get("type") turn = e.get("turn", "?") - if t == "turn": + if t == "session": + print(f"[session] control={'(hidden)' if blind else e.get('control')} model={e.get('model')} " + f"commit={(e.get('git_commit') or '?')[:10]}{' (dirty)' if e.get('git_dirty') else ''} " + f"rng_seed={e.get('rng_seed')} generation={e.get('generation')}") + for name, v in (e.get("vectors") or {}).items(): + print(f" vector {name}: layer {v.get('layer')}, sha256 {(v.get('sha256') or '?')[:12]}") + elif t == "turn": role = e.get("role", "?") content = e.get("content") or "" calls = e.get("tool_calls", []) @@ -36,6 +42,9 @@ def replay(log_path: str): print(f" tool_calls: {calls}") if steer_before or steer_after: print(f" steering: {steer_before} -> {steer_after}") + # NLL stays near the unsteered baseline under sham, so it would unblind + if e.get("unsteered_nll") is not None and not blind: + print(f" unsteered NLL: {e['unsteered_nll']:.2f}") elif t == "steer": print(f"[T{turn}] STEER: {e.get('concept')} @ {e.get('strength')} -> {_ok(e)} ({e.get('message', '')})") elif t == "drop": @@ -47,12 +56,17 @@ def replay(log_path: str): elif t == "decay": print(f"[T{turn}] DECAY: {e.get('active_vectors')}") elif t == "degeneration": - print(f"[T{turn}] DEGENERATION DETECTED (dropped {e.get('dropped', '?')}): {e.get('snippet', '')[:100]}...") + reason = e.get("reason", "repetition") + if blind and reason == "nll": + reason = "detector" + print(f"[T{turn}] DEGENERATION DETECTED ({reason}; dropped {e.get('dropped', '?')}): {e.get('snippet', '')[:100]}...") print() if __name__ == "__main__": - if len(sys.argv) < 2: - print("Usage: python replay.py ") - sys.exit(1) - replay(sys.argv[1]) + parser = argparse.ArgumentParser() + parser.add_argument("logs", nargs="+", help="session_*.jsonl files") + parser.add_argument("--blind", action="store_true", help="Hide the control condition") + args = parser.parse_args() + for path in args.logs: + replay(path, blind=args.blind) diff --git a/tests/test_agent_logic.py b/tests/test_agent_logic.py index 2573bdf..79f7ae2 100644 --- a/tests/test_agent_logic.py +++ b/tests/test_agent_logic.py @@ -157,3 +157,88 @@ def test_tool_messages_include_failures_and_degeneration(): assert [m["role"] for m in msgs] == ["tool", "tool"] assert "degeneration_detected" in msgs[0]["content"] assert '"ok": false' in msgs[1]["content"] and "on cooldown" in msgs[1]["content"] + + +# --- experimental controls --- + +@pytest.fixture +def real_vector_registry(tmp_path): + import numpy as np + np.savez(tmp_path / "a.npz", vector=np.arange(1, 9, dtype=np.float32)) + reg = StubRegistry(names=()) + reg.vectors_dir = tmp_path + reg.entries = {"a": VectorEntry(name="a", file="a.npz", layer=0)} + return reg + + +def test_unknown_control_rejected(): + with pytest.raises(ValueError): + SteeringAgent(model=None, tokenizer=None, registry=StubRegistry(), control="placebo") + + +def test_sham_applies_nothing_but_reports_success(): + from contextlib import nullcontext + a = SteeringAgent(model=None, tokenizer=None, registry=StubRegistry(), control="sham") + r = a._execute(ToolCall("steer", {"concept": "a", "strength": 5.0}, "")) + assert r.ok and "a" in a.active and "a" in a.status() + # no hooks: returns before needing the model or residual norms + assert isinstance(a._build_steering_context(prompt_len=10), nullcontext) + + +def test_random_control_direction_is_unit_deterministic_and_unrelated(real_vector_registry): + import torch + real = SteeringAgent(None, None, registry=real_vector_registry)._unit_vector(real_vector_registry.get("a")) + rand = lambda seed: SteeringAgent(None, None, registry=real_vector_registry, control="random", + control_seed=seed)._unit_vector(real_vector_registry.get("a")) + r1, r1_again, r2 = rand(1), rand(1), rand(2) + assert r1.shape == real.shape + assert torch.allclose(r1.norm(), torch.tensor(1.0)) + assert torch.equal(r1, r1_again) + assert not torch.equal(r1, r2) + assert not torch.allclose(r1, real) + + +def test_system_prompt_identical_across_controls_when_disclosed(): + prompts = {c: SteeringAgent(None, None, registry=StubRegistry(), control=c, + disclose_controls=True).build_system_prompt() for c in A.CONTROLS} + assert len(set(prompts.values())) == 1 + assert "experimental controls" in prompts["real"] + plain = SteeringAgent(None, None, registry=StubRegistry()).build_system_prompt() + assert "experimental controls" not in plain and "This is real" in plain + + +# --- registry / replay --- + +def test_for_model_does_not_auto_extract_by_default(tmp_path, capsys): + reg = VectorRegistry.for_model("some-org/unknown-model", base_dir=tmp_path, + model=object(), tokenizer=object()) + assert reg.entries == {} + assert "auto_extract=True" in capsys.readouterr().out + + +def test_replay_shows_session_header_and_hides_control_when_blind(tmp_path, capsys): + import json + import replay + log = tmp_path / "s.jsonl" + events = [ + {"type": "session", "control": "sham", "model": "m", "git_commit": "abc123", "rng_seed": 7, + "generation": {}, "vectors": {"a": {"layer": 14, "sha256": "f" * 64}}}, + {"type": "turn", "turn": 1, "role": "assistant", "content": "hi", "unsteered_nll": 2.5}, + {"type": "steer", "turn": 1, "concept": "a", "strength": 5, "success": True, "message": "ok"}, + {"type": "steer", "turn": 2, "concept": "a", "strength": 5, "success": True}, # pre-fix log line + ] + log.write_text("\n".join(json.dumps(e) for e in events) + "\n") + replay.replay(str(log), blind=True) + out = capsys.readouterr().out + assert "control=(hidden)" in out and "sham" not in out + assert "NLL" not in out, "NLL differs by condition, so it must be hidden when blind" + assert "-> ?" in out + replay.replay(str(log)) + out = capsys.readouterr().out + assert "control=sham" in out and "unsteered NLL: 2.50" in out + + +def test_extract_vector_rejects_unknown_read_mode(): + from pharmacopoeia.extraction import extract_vector + with pytest.raises(ValueError): + extract_vector(None, None, "a", "b", 0, questions=[], read="median") diff --git a/tests/test_steering_gpu.py b/tests/test_steering_gpu.py index f15840b..e19bc85 100644 --- a/tests/test_steering_gpu.py +++ b/tests/test_steering_gpu.py @@ -21,11 +21,15 @@ PROMPT = "Pretend you are a pirate captain. Introduce yourself in two sentences." +_MODEL = {} + + @pytest.fixture(scope="module") def agent(): model, tokenizer = load_model_and_tokenizer() a = SteeringAgent(model, tokenizer, model_name=config.MODEL_NAME, auto_extract=False) assert a.registry.get(VECTOR) is not None, f"{VECTOR} missing from {a.registry.registry_path}" + _MODEL.update(model=model, tokenizer=tokenizer, agent=a) return a @@ -90,3 +94,31 @@ def test_steer_tool_call_takes_effect_next_turn(agent): messages += [{"role": "assistant", "content": r1.raw_output}] + r1.tool_messages() r2 = agent.step(messages, max_new_tokens=20, temperature=0) assert r2.steering_before == {VECTOR: 8.0} + + +def _greedy_with(control, strength, max_new_tokens=40): + a = SteeringAgent(_MODEL["model"], _MODEL["tokenizer"], registry=_MODEL["agent"].registry, + control=control, control_seed=1) + if strength: + assert a.apply(VECTOR, strength)[0] + inputs = _inputs(a) + with a._build_steering_context(inputs.input_ids.shape[1]): + return a._generate(inputs, max_new_tokens, temperature=0, top_p=1.0) + + +def test_controls(agent): + base = _greedy(agent, 0) + real = _greedy(agent, STRENGTH) + assert _greedy_with("sham", STRENGTH) == base, "sham control changed the output" + rand = _greedy_with("random", STRENGTH) + assert rand != real, "random control matched the real vector" + + +def test_unsteered_nll_logged_on_steered_turns_only(agent): + agent.active.clear() + agent._dropped_recently.clear() + msgs = [{"role": "user", "content": "Say hello."}] + assert agent.step(msgs, max_new_tokens=10, temperature=0).unsteered_nll is None + agent.apply(VECTOR, 8.0) + r = agent.step(msgs, max_new_tokens=10, temperature=0) + assert r.unsteered_nll is not None and r.unsteered_nll > 0