diff --git a/AGENTS.md b/AGENTS.md index 28655a5..bacdd7e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -19,7 +19,13 @@ messages += [{"role": "assistant", "content": result.raw_output}] messages += result.tool_messages() # tool results back to the model ``` -The agent harness calls `agent.step(messages)` each turn instead of `model.generate()`, then appends the assistant output and `result.tool_messages()` (role `tool`) so the model sees whether its calls worked. If there are tool messages, give the model the next turn. Steering is transparent — the harness doesn't need to know about hooks or vectors. +Install it into the harness's environment with `pip install -e /path/to/modulo`. Keep the install editable so the default `vectors/` directory still resolves to the repo (or set `MODULO_VECTORS_DIR`). + +The agent harness calls `agent.step(messages)` each turn instead of `model.generate()`, then appends the assistant output and `result.tool_messages()` (role `tool`) so the model sees whether its calls worked. If there are tool messages, give the model the next turn. If `messages` starts with a system message, it is merged into the steering system prompt. Steering is transparent — the harness doesn't need to know about hooks or vectors. + +`agent.apply()` / `agent.drop()` called between steps behave like the model's own tool calls: an applied vector gets one full-strength steered turn, and a drop blocks re-applying for `COOLDOWN_TURNS` turns. + +Chat templates must accept `system` and `tool` roles (Qwen, Llama-3.x, recent Mistral do). Gemma-2 rejects both, so it isn't supported yet. ### 2. As a standalone demo diff --git a/agent/loop.py b/agent/loop.py index 030ca1b..f1d1405 100644 --- a/agent/loop.py +++ b/agent/loop.py @@ -71,7 +71,8 @@ def run_turn(agent: SteeringAgent, messages: list, logger: SessionLogger, turn: else: logger.log({"type": "tool_error", "turn": turn, "name": r.name, "message": r.message}) if result.steering_before != result.steering_after: - logger.log({"type": "decay", "turn": turn, "active_vectors": result.steering_after}) + # Steering changed this turn (steer/drop calls, decay, or expiry) + logger.log({"type": "steering_state", "turn": turn, "active_vectors": result.steering_after}) return result @@ -90,8 +91,10 @@ def main(): messages = [] if args.seed: messages.append({"role": "user", "content": args.seed}) + logger.log_turn(1, "user", args.seed) print(f"\n[seed] {args.seed}\n") + # Inputs (user, tool results, nudges) are logged under the turn they feed into turn = 0 while turn < args.max_turns and not kill.killed: if not messages or messages[-1]["role"] not in ("user", "tool"): @@ -105,7 +108,7 @@ def main(): if not user_input or user_input.lower() in ("quit", "exit"): break messages.append({"role": "user", "content": user_input}) - logger.log_turn(turn, "user", user_input) + logger.log_turn(turn + 1, "user", user_input) print(f"\n--- Turn {turn+1} ---") print(f"Steering status:\n{agent.status()}") @@ -119,12 +122,12 @@ def main(): # Tool results go straight back to the model, which gets the next turn for m in tool_messages: print(f"[tool] {m['content']}") - logger.log_turn(turn, "tool", m["content"]) + logger.log_turn(turn + 1, "tool", m["content"]) messages.extend(tool_messages) - elif args.auto: + elif args.auto and turn < args.max_turns and not kill.killed: # no nudge after the last turn nudge = "(You can continue experimenting with steer if you like, or just talk.)" messages.append({"role": "user", "content": nudge}) - logger.log_turn(turn, "user", nudge) + logger.log_turn(turn + 1, "user", nudge) print(f"\n[session ended after {turn} turns]") print(f"[log saved to {logger.log_path}]") diff --git a/config.py b/config.py index 2380756..c6b9642 100644 --- a/config.py +++ b/config.py @@ -9,22 +9,15 @@ # Qwen2.5-7B has 28 decoder layers (indices 0..27). Mid layers are good for steering. # CAA on Llama-2-7B (32 layers) used layer 13; for Qwen2.5-7B we start at layer 13. -NUM_LAYERS = 28 DEFAULT_STEER_LAYER = 13 LAYER_SWEEP = [10, 13, 16, 20] # tried if default doesn't shift behavior -# Steering -DEFAULT_MULTIPLIER = 2.0 -MULTIPLIER_SWEEP = [-2.0, -1.0, 0.0, 1.0, 2.0] - # Paths (resolved relative to project root) PROJECT_ROOT = Path(__file__).resolve().parent -DATA_DIR = PROJECT_ROOT / "data" VECTORS_DIR = PROJECT_ROOT / "vectors" LOGS_DIR = PROJECT_ROOT / "logs" -AGENT_DIR = PROJECT_ROOT / "agent" -for d in (DATA_DIR, VECTORS_DIR, LOGS_DIR, AGENT_DIR): +for d in (VECTORS_DIR, LOGS_DIR): d.mkdir(parents=True, exist_ok=True) # HF cache on the pod's large /workspace disk — override via env if needed. diff --git a/modulo/agent.py b/modulo/agent.py index 7468342..5606812 100644 --- a/modulo/agent.py +++ b/modulo/agent.py @@ -12,6 +12,7 @@ # agent.status() -> active vector status string """ import json +import math import re from collections import Counter from contextlib import nullcontext @@ -293,6 +294,9 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None, self.active: dict[str, ActiveVector] = {} self.turn = 0 self._dropped_recently: dict[str, int] = {} + # True between this turn's generation and decay(). apply()/drop() calls + # before generation (manual control) count as the end of the previous turn. + self._turn_generated = False self._unit_vectors: dict[str, torch.Tensor] = {} self._residual_norms: dict[int, float] = {} @@ -347,6 +351,7 @@ def apply(self, name: str, strength: float) -> tuple[bool, str]: return False, f"Too many active vectors (max {MAX_ACTIVE_VECTORS}). Drop one first." self.active[name] = ActiveVector( entry=entry, base_strength=strength, current_strength=strength, + turns_active=0 if self._turn_generated else 1, last_applied_turn=self.turn, ) sign = "+" if strength >= 0 else "" @@ -357,11 +362,12 @@ def drop(self, name: str) -> tuple[bool, str]: if name not in self.active: return False, f"'{name}' is not active." del self.active[name] - self._dropped_recently[name] = self.turn + self._dropped_recently[name] = self.turn if self._turn_generated else self.turn - 1 return True, f"Dropped '{name}'. On cooldown for the next {COOLDOWN_TURNS} turns." def decay(self): self.turn += 1 + self._turn_generated = False expired = [] for name, av in self.active.items(): av.turns_active += 1 @@ -377,7 +383,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() - self._unit_vectors[entry.name] = v / v.norm() + norm = v.norm() + if norm == 0 or not torch.isfinite(norm): + raise ValueError(f"Vector '{entry.name}' has zero or non-finite norm; re-extract it.") + self._unit_vectors[entry.name] = v / norm return self._unit_vectors[entry.name] def residual_norm(self, layer: int) -> float: @@ -426,7 +435,7 @@ def _execute(self, tc: ToolCall) -> ToolResult: strength = float(args.get("strength", 0)) except (TypeError, ValueError): return ToolResult(tc.name, args, False, f"'strength' must be a number, got {args.get('strength')!r}.") - if strength != strength: # NaN + if not math.isfinite(strength): return ToolResult(tc.name, args, False, "'strength' must be a finite number.") ok, msg = self.apply(concept, strength) elif tc.name == "drop": @@ -450,6 +459,20 @@ def _generate(self, inputs, max_new_tokens, temperature, top_p): ) return self.tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) + def _full_messages(self, messages: list[dict]) -> list[dict]: + sys_prompt = self.build_system_prompt() + if messages and messages[0].get("role") == "system": + # Merge the harness's own system prompt; many chat templates reject + # or mishandle a second system message. + sys_prompt = messages[0]["content"] + "\n\n" + sys_prompt + messages = messages[1:] + return [{"role": "system", "content": sys_prompt}] + messages + + def _encode(self, messages: list[dict]): + text = self.tokenizer.apply_chat_template( + self._full_messages(messages), tokenize=False, add_generation_prompt=True) + return self.tokenizer([text], return_tensors="pt").to(self.model.device) + def step(self, messages: list[dict], max_new_tokens: int = 300, temperature: float = 0.7, top_p: float = 0.9) -> StepResult: """Run one turn: build prompt, generate (steered), parse tool calls, apply them, decay. @@ -460,16 +483,14 @@ def step(self, messages: list[dict], max_new_tokens: int = 300, Returns: StepResult with prose, raw output, tool calls/results, and steering state """ - sys_prompt = self.build_system_prompt() - full_messages = [{"role": "system", "content": sys_prompt}] + messages - text = self.tokenizer.apply_chat_template(full_messages, tokenize=False, add_generation_prompt=True) - inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device) + inputs = self._encode(messages) prompt_len = inputs.input_ids.shape[1] 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) + self._turn_generated = True # Degeneration check (only meaningful if steering was applied) degenerate_output = None @@ -477,7 +498,8 @@ def step(self, messages: list[dict], max_new_tokens: int = 300, degenerate_output = raw_output for name in list(self.active.keys()): self.drop(name) - raw_output = self._generate(inputs, max_new_tokens, temperature, top_p) + # Re-encode so the regenerated turn's system prompt shows the drops + raw_output = self._generate(self._encode(messages), max_new_tokens, temperature, top_p) tool_calls = parse_tool_calls(raw_output) prose = strip_tool_calls(raw_output) diff --git a/modulo/extraction.py b/modulo/extraction.py index bd33b4d..94ee873 100644 --- a/modulo/extraction.py +++ b/modulo/extraction.py @@ -85,9 +85,11 @@ def hook(module, input, output): activation = act[0, -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 diff --git a/modulo/registry.py b/modulo/registry.py index a654091..f03d7fc 100644 --- a/modulo/registry.py +++ b/modulo/registry.py @@ -128,14 +128,14 @@ def load_vector(self, vectors_dir: Path) -> torch.Tensor: """Load the actual vector tensor for this entry's layer.""" path = vectors_dir / self.file if self.format == "per_layer": - data = torch.load(path, weights_only=True) + data = torch.load(path, map_location="cpu", weights_only=True) return data[self.layer].to(torch.bfloat16) elif self.file.endswith(".npz"): import numpy as np data = np.load(path) return torch.tensor(data["vector"], dtype=torch.bfloat16) elif self.file.endswith(".pt"): - data = torch.load(path, weights_only=True) + data = torch.load(path, map_location="cpu", weights_only=True) return data.to(torch.bfloat16) else: raise ValueError(f"Unknown vector format: {self.file}") diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..2b5b7e6 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = ["setuptools>=64"] +build-backend = "setuptools.build_meta" + +[project] +name = "modulo" +version = "0.1.0" +description = "Self-steering for local language models: let a model steer its own activations." +requires-python = ">=3.10" +dependencies = [ + "torch>=2.4", + "transformers>=4.46,<5", + "accelerate>=1.0", + "huggingface_hub", + "numpy", + "pyyaml", +] + +[project.optional-dependencies] +# Extra packages used by the repo's scripts (CAA extraction, quantization) and tests +scripts = ["steering-vectors>=0.12", "datasets>=3.0", "scipy", "bitsandbytes"] +test = ["pytest"] + +[tool.setuptools] +# Only the library is installable; the demo, scripts and config.py stay repo-local. +# Install editable (pip install -e .) so the default vectors/ dir resolves to this repo. +packages = ["modulo"] diff --git a/replay.py b/replay.py index a077cc9..ac59e31 100644 --- a/replay.py +++ b/replay.py @@ -44,8 +44,8 @@ def replay(log_path: str): print(f"[T{turn}] REPORT: {e.get('concept')} worked={e.get('worked')} — {e.get('notes')}") elif t == "tool_error": print(f"[T{turn}] TOOL ERROR: {e.get('name')} — {e.get('message', '')}") - elif t == "decay": - print(f"[T{turn}] DECAY: {e.get('active_vectors')}") + elif t in ("steering_state", "decay"): # "decay" in older logs + print(f"[T{turn}] STEERING: {e.get('active_vectors')}") elif t == "degeneration": print(f"[T{turn}] DEGENERATION DETECTED (dropped {e.get('dropped', '?')}): {e.get('snippet', '')[:100]}...") print() diff --git a/setup.sh b/setup.sh index a023107..feec90d 100755 --- a/setup.sh +++ b/setup.sh @@ -30,7 +30,7 @@ if ! pip install --break-system-packages --quiet -r requirements.txt; then fi # 2. Create directories -mkdir -p "$VEC_DIR" data logs agent +mkdir -p "$VEC_DIR" logs # 3. Download the assistant-axis steering vector (the one confirmed-working vector) echo "" diff --git a/tests/test_agent_logic.py b/tests/test_agent_logic.py index 94fa70a..7a0a545 100644 --- a/tests/test_agent_logic.py +++ b/tests/test_agent_logic.py @@ -2,7 +2,10 @@ No model is loaded; SteeringAgent runs against a stub registry. """ +from contextlib import nullcontext + import pytest +import torch import modulo.agent as A from modulo.agent import (COOLDOWN_TURNS, DECAY_FACTOR, SteeringAgent, ToolCall, @@ -85,7 +88,13 @@ def test_strength_clamped_to_entry_and_global_cap(agent): assert agent.active["a"].current_strength == cap +def in_step(agent): + """Put the agent where step() runs tool calls: after generation, before decay.""" + agent._turn_generated = True + + def test_full_strength_on_first_turn_then_decays(agent): + in_step(agent) agent.apply("a", 5.0) agent.decay() # end of the turn steer was called; next turn runs at full strength assert agent.active["a"].current_strength == 5.0 @@ -94,6 +103,7 @@ def test_full_strength_on_first_turn_then_decays(agent): def test_cooldown_blocks_exactly_cooldown_turns(agent): + in_step(agent) agent.apply("a", 1.0) assert agent.drop("a")[0] agent.decay() @@ -106,6 +116,31 @@ def test_cooldown_blocks_exactly_cooldown_turns(agent): assert blocked == COOLDOWN_TURNS +def test_manual_apply_between_steps_gets_one_full_strength_turn(agent): + agent.apply("a", 5.0) # harness call before step() + assert "turn 1" in agent.status() + assert agent.active["a"].current_strength == 5.0 # strength used for this turn's generation + in_step(agent) + agent.decay() + assert agent.active["a"].current_strength == pytest.approx(5.0 * DECAY_FACTOR) + + +def test_manual_drop_between_steps_blocks_exactly_cooldown_turns(agent): + in_step(agent) + agent.apply("a", 1.0) + agent.decay() + assert agent.drop("a")[0] # harness call before step() + assert f"{COOLDOWN_TURNS} more turns" in agent.status() + blocked = 0 + for _ in range(COOLDOWN_TURNS + 3): + in_step(agent) + if agent.apply("a", 1.0)[0]: + break + blocked += 1 + agent.decay() + assert blocked == COOLDOWN_TURNS + + def test_natural_expiry_has_no_cooldown(agent): agent.apply("b", 0.2) for _ in range(10): @@ -127,6 +162,8 @@ def test_max_active_vectors(agent): @pytest.mark.parametrize("args", [ {"concept": "a", "strength": "high"}, {"concept": "a", "strength": "nan"}, + {"concept": "a", "strength": "inf"}, + {"concept": "a", "strength": "-inf"}, {"concept": ["a"], "strength": 1}, {"strength": 1}, {"concept": "nope", "strength": 1}, @@ -157,3 +194,37 @@ 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"] + + +# --- prompt building --- + +def test_harness_system_prompt_is_merged(agent): + full = agent._full_messages([{"role": "system", "content": "HARNESS RULES"}, + {"role": "user", "content": "hi"}]) + assert [m["role"] for m in full] == ["system", "user"] + assert full[0]["content"].startswith("HARNESS RULES\n\n") + assert "self-steering experiment" in full[0]["content"] + + +def test_degenerate_turn_regenerates_with_updated_prompt(agent, monkeypatch): + """The unsteered retry must be built after the drops, so its prompt doesn't list them as active.""" + class FakeInputs: + def __init__(self, prompt): + self.prompt = prompt + self.input_ids = torch.zeros(1, 1) + + prompts, outputs = [], ["loop " * 50, "Fine now."] + + def fake_generate(inputs, *args): + prompts.append(inputs.prompt) + return outputs.pop(0) + + monkeypatch.setattr(agent, "_encode", lambda msgs: FakeInputs(agent._full_messages(msgs)[0]["content"])) + monkeypatch.setattr(agent, "_build_steering_context", lambda prompt_len: nullcontext()) + monkeypatch.setattr(agent, "_generate", fake_generate) + + agent.apply("a", 5.0) + result = agent.step([{"role": "user", "content": "hi"}]) + assert result.degenerate_output is not None and result.raw_output == "Fine now." + assert "a: +5.00" in prompts[0] + assert "a: +5.00" not in prompts[1] and "a: on cooldown" in prompts[1]