weightspace 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- weightspace-0.1.0/.github/workflows/ci.yml +14 -0
- weightspace-0.1.0/.github/workflows/publish.yml +30 -0
- weightspace-0.1.0/.gitignore +9 -0
- weightspace-0.1.0/PKG-INFO +181 -0
- weightspace-0.1.0/README.md +148 -0
- weightspace-0.1.0/examples/fix_quadrants.py +86 -0
- weightspace-0.1.0/examples/hf_reduce_loops.py +106 -0
- weightspace-0.1.0/examples/read_which_task.py +56 -0
- weightspace-0.1.0/pyproject.toml +49 -0
- weightspace-0.1.0/src/weightspace/__init__.py +45 -0
- weightspace-0.1.0/src/weightspace/adapters.py +112 -0
- weightspace-0.1.0/src/weightspace/agent.py +436 -0
- weightspace-0.1.0/src/weightspace/cli.py +73 -0
- weightspace-0.1.0/src/weightspace/llm/__init__.py +35 -0
- weightspace-0.1.0/src/weightspace/llm/anthropic.py +125 -0
- weightspace-0.1.0/src/weightspace/llm/base.py +71 -0
- weightspace-0.1.0/src/weightspace/llm/mock.py +38 -0
- weightspace-0.1.0/src/weightspace/llm/openai.py +108 -0
- weightspace-0.1.0/src/weightspace/probes.py +60 -0
- weightspace-0.1.0/src/weightspace/prompts.py +43 -0
- weightspace-0.1.0/src/weightspace/render.py +151 -0
- weightspace-0.1.0/src/weightspace/sandbox.py +68 -0
- weightspace-0.1.0/src/weightspace/space.py +319 -0
- weightspace-0.1.0/src/weightspace/tasks.py +57 -0
- weightspace-0.1.0/tests/conftest.py +30 -0
- weightspace-0.1.0/tests/test_agent.py +136 -0
- weightspace-0.1.0/tests/test_llm.py +148 -0
- weightspace-0.1.0/tests/test_render.py +34 -0
- weightspace-0.1.0/tests/test_space.py +75 -0
|
@@ -0,0 +1,14 @@
|
|
|
1
|
+
name: ci
|
|
2
|
+
on: [push, pull_request]
|
|
3
|
+
jobs:
|
|
4
|
+
test:
|
|
5
|
+
runs-on: ubuntu-latest
|
|
6
|
+
steps:
|
|
7
|
+
- uses: actions/checkout@v4
|
|
8
|
+
- uses: actions/setup-python@v5
|
|
9
|
+
with:
|
|
10
|
+
python-version: "3.11"
|
|
11
|
+
- run: pip install torch --index-url https://download.pytorch.org/whl/cpu
|
|
12
|
+
- run: pip install -e ".[all,dev]"
|
|
13
|
+
- run: ruff check . && ruff format --check .
|
|
14
|
+
- run: pytest -q
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
name: publish
|
|
2
|
+
on:
|
|
3
|
+
push:
|
|
4
|
+
tags: ["v*"]
|
|
5
|
+
workflow_dispatch:
|
|
6
|
+
jobs:
|
|
7
|
+
build:
|
|
8
|
+
runs-on: ubuntu-latest
|
|
9
|
+
steps:
|
|
10
|
+
- uses: actions/checkout@v4
|
|
11
|
+
- uses: actions/setup-python@v5
|
|
12
|
+
with:
|
|
13
|
+
python-version: "3.11"
|
|
14
|
+
- run: pip install build && python -m build
|
|
15
|
+
- uses: actions/upload-artifact@v4
|
|
16
|
+
with:
|
|
17
|
+
name: dist
|
|
18
|
+
path: dist/
|
|
19
|
+
pypi:
|
|
20
|
+
needs: build
|
|
21
|
+
runs-on: ubuntu-latest
|
|
22
|
+
environment: pypi
|
|
23
|
+
permissions:
|
|
24
|
+
id-token: write
|
|
25
|
+
steps:
|
|
26
|
+
- uses: actions/download-artifact@v4
|
|
27
|
+
with:
|
|
28
|
+
name: dist
|
|
29
|
+
path: dist/
|
|
30
|
+
- uses: pypa/gh-action-pypi-publish@release/v1
|
|
@@ -0,0 +1,181 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: weightspace
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Weight-space programming: let an LLM read and write neural network weights directly.
|
|
5
|
+
Project-URL: Homepage, https://github.com/leofmarciano/weightspace
|
|
6
|
+
Project-URL: Source, https://github.com/leofmarciano/weightspace
|
|
7
|
+
Keywords: interpretability,llm,model-editing,pytorch,weights
|
|
8
|
+
Classifier: Intended Audience :: Science/Research
|
|
9
|
+
Classifier: Programming Language :: Python :: 3
|
|
10
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
11
|
+
Requires-Python: >=3.10
|
|
12
|
+
Requires-Dist: numpy>=1.24
|
|
13
|
+
Provides-Extra: all
|
|
14
|
+
Requires-Dist: anthropic>=0.40; extra == 'all'
|
|
15
|
+
Requires-Dist: openai>=1.40; extra == 'all'
|
|
16
|
+
Requires-Dist: safetensors>=0.4; extra == 'all'
|
|
17
|
+
Requires-Dist: torch>=2.1; extra == 'all'
|
|
18
|
+
Requires-Dist: transformers>=4.40; extra == 'all'
|
|
19
|
+
Provides-Extra: anthropic
|
|
20
|
+
Requires-Dist: anthropic>=0.40; extra == 'anthropic'
|
|
21
|
+
Provides-Extra: dev
|
|
22
|
+
Requires-Dist: pytest>=8; extra == 'dev'
|
|
23
|
+
Requires-Dist: ruff>=0.5; extra == 'dev'
|
|
24
|
+
Provides-Extra: hf
|
|
25
|
+
Requires-Dist: safetensors>=0.4; extra == 'hf'
|
|
26
|
+
Requires-Dist: torch>=2.1; extra == 'hf'
|
|
27
|
+
Requires-Dist: transformers>=4.40; extra == 'hf'
|
|
28
|
+
Provides-Extra: openai
|
|
29
|
+
Requires-Dist: openai>=1.40; extra == 'openai'
|
|
30
|
+
Provides-Extra: torch
|
|
31
|
+
Requires-Dist: torch>=2.1; extra == 'torch'
|
|
32
|
+
Description-Content-Type: text/markdown
|
|
33
|
+
|
|
34
|
+
# weightspace
|
|
35
|
+
|
|
36
|
+
A small Python framework for **weight-space programming**: an LLM reads and edits a neural network's weights directly, by reasoning about what each part computes, instead of training it.
|
|
37
|
+
|
|
38
|
+
Based on [Sparks of Weight-Space Programming](https://rikiyat.github.io/blog/weight-space-programming/) (Takehi, Teoh, Zhang, Khattab, 2026). This repo is the harness, not the paper's results: how well it works depends on the LLM you plug in.
|
|
39
|
+
|
|
40
|
+
## What you get
|
|
41
|
+
|
|
42
|
+
- **Any model.** PyTorch modules, Hugging Face models, or `.safetensors` / `.npz` / `.pt` files. The LLM only ever sees tensors, never the model object.
|
|
43
|
+
- **Four tasks.** `read` (what does this network do?), `fix` (find and fix the cause of a failure), `write` (program a network from zero), `improve` (make it better at a goal).
|
|
44
|
+
- **Tools for the LLM.** List, describe (norms, SVD) and read tensor slices; a persistent Python sandbox over the weights; edits, checkpoints and rollbacks.
|
|
45
|
+
- **Budgeted forward passes.** You decide how the network may be run (scores, rollouts, generations, attention maps only...) by writing a `probe`, and how many times.
|
|
46
|
+
- **Edit tracking.** Every run reports exactly how many weights changed, and exports a sparse patch.
|
|
47
|
+
- **Text mode.** For small networks: dump every weight into the prompt and get one reply back, no tools (like the paper's first experiments).
|
|
48
|
+
- **Providers.** OpenAI-compatible APIs (OpenAI, OpenRouter, vLLM, Ollama) and Anthropic, or your own `LLM` class.
|
|
49
|
+
|
|
50
|
+
## Install
|
|
51
|
+
|
|
52
|
+
```bash
|
|
53
|
+
pip install -e ".[all]" # or pick: .[torch] .[hf] .[openai] .[anthropic]
|
|
54
|
+
export OPENAI_API_KEY=... # or ANTHROPIC_API_KEY / OPENROUTER_API_KEY
|
|
55
|
+
```
|
|
56
|
+
|
|
57
|
+
## Quick start
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
import torch
|
|
61
|
+
import weightspace as wsp
|
|
62
|
+
|
|
63
|
+
model = ... # any torch.nn.Module
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
@wsp.probe
|
|
67
|
+
def accuracy():
|
|
68
|
+
"""Per-class accuracy on a held-out batch."""
|
|
69
|
+
with torch.no_grad():
|
|
70
|
+
pred = model(x_val).argmax(-1)
|
|
71
|
+
return {c: float((pred[y_val == c] == c).float().mean()) for c in range(4)}
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
res = wsp.fix(
|
|
75
|
+
model,
|
|
76
|
+
"The classifier never predicts class 3. Find out why and fix it.",
|
|
77
|
+
llm="openai:gpt-5", # or "anthropic:claude-sonnet-4-5", "openrouter:..."
|
|
78
|
+
context="Linear(2->16) -> ReLU -> Linear(16->4); input is a 2D point, output is its quadrant.",
|
|
79
|
+
probes=[accuracy],
|
|
80
|
+
forward_budget=3,
|
|
81
|
+
)
|
|
82
|
+
|
|
83
|
+
print(res) # answer, confidence, forward passes used, which weights changed
|
|
84
|
+
res.save("runs/fix") # transcript.jsonl, result.json, patch.npz
|
|
85
|
+
```
|
|
86
|
+
|
|
87
|
+
The model is edited in place. The other tasks look the same:
|
|
88
|
+
|
|
89
|
+
```python
|
|
90
|
+
wsp.read(model, "Which game was this policy trained on?", llm=..., context=GRID_AND_ACTIONS)
|
|
91
|
+
wsp.read(model, "What does this network compute?", llm=..., text_mode=True) # weights in the prompt
|
|
92
|
+
wsp.write(model, "Play Breakout: move the paddle under the ball.", llm=..., context=GAME_RULES) # zeroes first
|
|
93
|
+
wsp.improve(hf_model, "Stop the reasoning loops.", llm=..., probes=[generate], forward_budget=25)
|
|
94
|
+
```
|
|
95
|
+
|
|
96
|
+
## Thinking on or off
|
|
97
|
+
|
|
98
|
+
`thinking` sets how much the editing LLM reasons before it acts. Use `thinking=False` to have it edit the weights with no thinking at all:
|
|
99
|
+
|
|
100
|
+
```python
|
|
101
|
+
wsp.fix(model, "...", llm="openai:gpt-5.1", thinking=False) # or "low" / "medium" / "high"; None = provider default
|
|
102
|
+
```
|
|
103
|
+
|
|
104
|
+
With thinking off, the provider's reasoning is disabled and the LLM is told to go straight to tool calls and edits. How that reaches each provider:
|
|
105
|
+
|
|
106
|
+
| Backend | off | on |
|
|
107
|
+
|---|---|---|
|
|
108
|
+
| OpenAI | `reasoning_effort="none"` (`OpenAIChat(..., off_effort="minimal")` for gpt-5) | `reasoning_effort=<level>` |
|
|
109
|
+
| OpenRouter | `reasoning.enabled=false` | `reasoning.effort=<level>` |
|
|
110
|
+
| vLLM / SGLang (Qwen3 etc.) | `OpenAIChat(..., thinking_style="chat_template")` → `enable_thinking=false` | `enable_thinking=true` |
|
|
111
|
+
| Anthropic | `thinking.type="disabled"` | extended thinking with a 2k / 8k / 16k token budget |
|
|
112
|
+
|
|
113
|
+
CLI: `--no-thinking` or `--thinking low|medium|high`. The setting is saved in `result.json`, so you can compare runs with and without thinking.
|
|
114
|
+
|
|
115
|
+
## Probes
|
|
116
|
+
|
|
117
|
+
A probe is any function. It runs after the current weights are synced into the model, and whatever it returns is what the LLM sees. Its arguments become the tool's schema.
|
|
118
|
+
|
|
119
|
+
```python
|
|
120
|
+
@wsp.probe(cost=1)
|
|
121
|
+
def generate(prompts: list, max_new_tokens: int = 256):
|
|
122
|
+
"""Generate answers for up to 8 prompts and report repetition."""
|
|
123
|
+
...
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
Want the LLM to see only attention patterns, not outputs? Write a probe that returns only attention patterns. That's how you set what the LLM can observe.
|
|
127
|
+
|
|
128
|
+
## Lower level
|
|
129
|
+
|
|
130
|
+
```python
|
|
131
|
+
ws = wsp.WeightSpace.from_hf("Qwen/Qwen2.5-0.5B-Instruct") # or from_torch / from_file / from_numpy
|
|
132
|
+
ws.W["model.layers.15.mlp.down_proj.weight"][:, :64] -= ... # tracked in-place edit
|
|
133
|
+
print(ws.diff()) # what changed vs. the original
|
|
134
|
+
ws.checkpoint("a")
|
|
135
|
+
ws.restore("a")
|
|
136
|
+
ws.revert()
|
|
137
|
+
ws.save_patch("fix.npz")
|
|
138
|
+
other.apply_patch("fix.npz")
|
|
139
|
+
|
|
140
|
+
s = wsp.Session(
|
|
141
|
+
ws,
|
|
142
|
+
llm,
|
|
143
|
+
mode="improve",
|
|
144
|
+
probes=[...],
|
|
145
|
+
forward_budget=50,
|
|
146
|
+
tools=[my_tool],
|
|
147
|
+
namespace={"tok": tokenizer},
|
|
148
|
+
max_steps=80,
|
|
149
|
+
verbose=True,
|
|
150
|
+
)
|
|
151
|
+
res = s.run("...")
|
|
152
|
+
```
|
|
153
|
+
|
|
154
|
+
Custom tools: `wsp.tool("name", "description", x={"type": "number"})(fn)`.
|
|
155
|
+
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script.
|
|
156
|
+
|
|
157
|
+
## CLI
|
|
158
|
+
|
|
159
|
+
```bash
|
|
160
|
+
weightspace inspect model.safetensors --describe
|
|
161
|
+
weightspace read model.safetensors "What was this trained to do?" --llm openai:gpt-5 --context @notes.md
|
|
162
|
+
weightspace fix model.safetensors "..." --llm anthropic:claude-sonnet-4-5 --save-weights fixed.safetensors --out runs/1
|
|
163
|
+
```
|
|
164
|
+
|
|
165
|
+
## Examples
|
|
166
|
+
|
|
167
|
+
- `examples/fix_quadrants.py`: fix a sabotaged classifier (runs offline with a scripted LLM if you omit `--llm`)
|
|
168
|
+
- `examples/read_which_task.py`: name the boolean function an MLP learned, from its weights in the prompt
|
|
169
|
+
- `examples/hf_reduce_loops.py`: edit a Hugging Face LM to stop repetition loops, with generation and attention probes
|
|
170
|
+
|
|
171
|
+
## Rules and caveats
|
|
172
|
+
|
|
173
|
+
The system prompt tells the LLM not to train, fit weights to data, or run the network inside the sandbox; the only way to run it is the budgeted probes. The sandbox never holds the model object, but it is not a security boundary: code runs in your process, and an LLM could still multiply matrices by hand. Check transcripts if the rules matter for your experiment, as the paper's authors did.
|
|
174
|
+
|
|
175
|
+
## Development
|
|
176
|
+
|
|
177
|
+
```bash
|
|
178
|
+
pip install -e ".[all,dev]"
|
|
179
|
+
pytest
|
|
180
|
+
ruff check . && ruff format --check .
|
|
181
|
+
```
|
|
@@ -0,0 +1,148 @@
|
|
|
1
|
+
# weightspace
|
|
2
|
+
|
|
3
|
+
A small Python framework for **weight-space programming**: an LLM reads and edits a neural network's weights directly, by reasoning about what each part computes, instead of training it.
|
|
4
|
+
|
|
5
|
+
Based on [Sparks of Weight-Space Programming](https://rikiyat.github.io/blog/weight-space-programming/) (Takehi, Teoh, Zhang, Khattab, 2026). This repo is the harness, not the paper's results: how well it works depends on the LLM you plug in.
|
|
6
|
+
|
|
7
|
+
## What you get
|
|
8
|
+
|
|
9
|
+
- **Any model.** PyTorch modules, Hugging Face models, or `.safetensors` / `.npz` / `.pt` files. The LLM only ever sees tensors, never the model object.
|
|
10
|
+
- **Four tasks.** `read` (what does this network do?), `fix` (find and fix the cause of a failure), `write` (program a network from zero), `improve` (make it better at a goal).
|
|
11
|
+
- **Tools for the LLM.** List, describe (norms, SVD) and read tensor slices; a persistent Python sandbox over the weights; edits, checkpoints and rollbacks.
|
|
12
|
+
- **Budgeted forward passes.** You decide how the network may be run (scores, rollouts, generations, attention maps only...) by writing a `probe`, and how many times.
|
|
13
|
+
- **Edit tracking.** Every run reports exactly how many weights changed, and exports a sparse patch.
|
|
14
|
+
- **Text mode.** For small networks: dump every weight into the prompt and get one reply back, no tools (like the paper's first experiments).
|
|
15
|
+
- **Providers.** OpenAI-compatible APIs (OpenAI, OpenRouter, vLLM, Ollama) and Anthropic, or your own `LLM` class.
|
|
16
|
+
|
|
17
|
+
## Install
|
|
18
|
+
|
|
19
|
+
```bash
|
|
20
|
+
pip install -e ".[all]" # or pick: .[torch] .[hf] .[openai] .[anthropic]
|
|
21
|
+
export OPENAI_API_KEY=... # or ANTHROPIC_API_KEY / OPENROUTER_API_KEY
|
|
22
|
+
```
|
|
23
|
+
|
|
24
|
+
## Quick start
|
|
25
|
+
|
|
26
|
+
```python
|
|
27
|
+
import torch
|
|
28
|
+
import weightspace as wsp
|
|
29
|
+
|
|
30
|
+
model = ... # any torch.nn.Module
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
@wsp.probe
|
|
34
|
+
def accuracy():
|
|
35
|
+
"""Per-class accuracy on a held-out batch."""
|
|
36
|
+
with torch.no_grad():
|
|
37
|
+
pred = model(x_val).argmax(-1)
|
|
38
|
+
return {c: float((pred[y_val == c] == c).float().mean()) for c in range(4)}
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
res = wsp.fix(
|
|
42
|
+
model,
|
|
43
|
+
"The classifier never predicts class 3. Find out why and fix it.",
|
|
44
|
+
llm="openai:gpt-5", # or "anthropic:claude-sonnet-4-5", "openrouter:..."
|
|
45
|
+
context="Linear(2->16) -> ReLU -> Linear(16->4); input is a 2D point, output is its quadrant.",
|
|
46
|
+
probes=[accuracy],
|
|
47
|
+
forward_budget=3,
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
print(res) # answer, confidence, forward passes used, which weights changed
|
|
51
|
+
res.save("runs/fix") # transcript.jsonl, result.json, patch.npz
|
|
52
|
+
```
|
|
53
|
+
|
|
54
|
+
The model is edited in place. The other tasks look the same:
|
|
55
|
+
|
|
56
|
+
```python
|
|
57
|
+
wsp.read(model, "Which game was this policy trained on?", llm=..., context=GRID_AND_ACTIONS)
|
|
58
|
+
wsp.read(model, "What does this network compute?", llm=..., text_mode=True) # weights in the prompt
|
|
59
|
+
wsp.write(model, "Play Breakout: move the paddle under the ball.", llm=..., context=GAME_RULES) # zeroes first
|
|
60
|
+
wsp.improve(hf_model, "Stop the reasoning loops.", llm=..., probes=[generate], forward_budget=25)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
## Thinking on or off
|
|
64
|
+
|
|
65
|
+
`thinking` sets how much the editing LLM reasons before it acts. Use `thinking=False` to have it edit the weights with no thinking at all:
|
|
66
|
+
|
|
67
|
+
```python
|
|
68
|
+
wsp.fix(model, "...", llm="openai:gpt-5.1", thinking=False) # or "low" / "medium" / "high"; None = provider default
|
|
69
|
+
```
|
|
70
|
+
|
|
71
|
+
With thinking off, the provider's reasoning is disabled and the LLM is told to go straight to tool calls and edits. How that reaches each provider:
|
|
72
|
+
|
|
73
|
+
| Backend | off | on |
|
|
74
|
+
|---|---|---|
|
|
75
|
+
| OpenAI | `reasoning_effort="none"` (`OpenAIChat(..., off_effort="minimal")` for gpt-5) | `reasoning_effort=<level>` |
|
|
76
|
+
| OpenRouter | `reasoning.enabled=false` | `reasoning.effort=<level>` |
|
|
77
|
+
| vLLM / SGLang (Qwen3 etc.) | `OpenAIChat(..., thinking_style="chat_template")` → `enable_thinking=false` | `enable_thinking=true` |
|
|
78
|
+
| Anthropic | `thinking.type="disabled"` | extended thinking with a 2k / 8k / 16k token budget |
|
|
79
|
+
|
|
80
|
+
CLI: `--no-thinking` or `--thinking low|medium|high`. The setting is saved in `result.json`, so you can compare runs with and without thinking.
|
|
81
|
+
|
|
82
|
+
## Probes
|
|
83
|
+
|
|
84
|
+
A probe is any function. It runs after the current weights are synced into the model, and whatever it returns is what the LLM sees. Its arguments become the tool's schema.
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
@wsp.probe(cost=1)
|
|
88
|
+
def generate(prompts: list, max_new_tokens: int = 256):
|
|
89
|
+
"""Generate answers for up to 8 prompts and report repetition."""
|
|
90
|
+
...
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
Want the LLM to see only attention patterns, not outputs? Write a probe that returns only attention patterns. That's how you set what the LLM can observe.
|
|
94
|
+
|
|
95
|
+
## Lower level
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
ws = wsp.WeightSpace.from_hf("Qwen/Qwen2.5-0.5B-Instruct") # or from_torch / from_file / from_numpy
|
|
99
|
+
ws.W["model.layers.15.mlp.down_proj.weight"][:, :64] -= ... # tracked in-place edit
|
|
100
|
+
print(ws.diff()) # what changed vs. the original
|
|
101
|
+
ws.checkpoint("a")
|
|
102
|
+
ws.restore("a")
|
|
103
|
+
ws.revert()
|
|
104
|
+
ws.save_patch("fix.npz")
|
|
105
|
+
other.apply_patch("fix.npz")
|
|
106
|
+
|
|
107
|
+
s = wsp.Session(
|
|
108
|
+
ws,
|
|
109
|
+
llm,
|
|
110
|
+
mode="improve",
|
|
111
|
+
probes=[...],
|
|
112
|
+
forward_budget=50,
|
|
113
|
+
tools=[my_tool],
|
|
114
|
+
namespace={"tok": tokenizer},
|
|
115
|
+
max_steps=80,
|
|
116
|
+
verbose=True,
|
|
117
|
+
)
|
|
118
|
+
res = s.run("...")
|
|
119
|
+
```
|
|
120
|
+
|
|
121
|
+
Custom tools: `wsp.tool("name", "description", x={"type": "number"})(fn)`.
|
|
122
|
+
Offline runs and tests: `wsp.ScriptedLLM([...])` replays a fixed script.
|
|
123
|
+
|
|
124
|
+
## CLI
|
|
125
|
+
|
|
126
|
+
```bash
|
|
127
|
+
weightspace inspect model.safetensors --describe
|
|
128
|
+
weightspace read model.safetensors "What was this trained to do?" --llm openai:gpt-5 --context @notes.md
|
|
129
|
+
weightspace fix model.safetensors "..." --llm anthropic:claude-sonnet-4-5 --save-weights fixed.safetensors --out runs/1
|
|
130
|
+
```
|
|
131
|
+
|
|
132
|
+
## Examples
|
|
133
|
+
|
|
134
|
+
- `examples/fix_quadrants.py`: fix a sabotaged classifier (runs offline with a scripted LLM if you omit `--llm`)
|
|
135
|
+
- `examples/read_which_task.py`: name the boolean function an MLP learned, from its weights in the prompt
|
|
136
|
+
- `examples/hf_reduce_loops.py`: edit a Hugging Face LM to stop repetition loops, with generation and attention probes
|
|
137
|
+
|
|
138
|
+
## Rules and caveats
|
|
139
|
+
|
|
140
|
+
The system prompt tells the LLM not to train, fit weights to data, or run the network inside the sandbox; the only way to run it is the budgeted probes. The sandbox never holds the model object, but it is not a security boundary: code runs in your process, and an LLM could still multiply matrices by hand. Check transcripts if the rules matter for your experiment, as the paper's authors did.
|
|
141
|
+
|
|
142
|
+
## Development
|
|
143
|
+
|
|
144
|
+
```bash
|
|
145
|
+
pip install -e ".[all,dev]"
|
|
146
|
+
pytest
|
|
147
|
+
ruff check . && ruff format --check .
|
|
148
|
+
```
|
|
@@ -0,0 +1,86 @@
|
|
|
1
|
+
"""Fix a sabotaged classifier by editing its weights.
|
|
2
|
+
|
|
3
|
+
A small MLP learns which quadrant a 2D point is in. We then break it so it
|
|
4
|
+
never predicts class 3, and ask an LLM to find and fix the cause, with a
|
|
5
|
+
budget of 3 forward passes.
|
|
6
|
+
|
|
7
|
+
python examples/fix_quadrants.py --llm openai:gpt-5
|
|
8
|
+
python examples/fix_quadrants.py --llm anthropic:claude-sonnet-4-5
|
|
9
|
+
python examples/fix_quadrants.py # scripted LLM, no API key needed
|
|
10
|
+
"""
|
|
11
|
+
|
|
12
|
+
import argparse
|
|
13
|
+
|
|
14
|
+
import torch
|
|
15
|
+
|
|
16
|
+
import weightspace as wsp
|
|
17
|
+
from weightspace.llm import ScriptedLLM, call, reply
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def train():
|
|
21
|
+
torch.manual_seed(0)
|
|
22
|
+
x = torch.randn(4096, 2)
|
|
23
|
+
y = (x[:, 0] > 0).long() * 2 + (x[:, 1] > 0).long()
|
|
24
|
+
model = torch.nn.Sequential(torch.nn.Linear(2, 16), torch.nn.ReLU(), torch.nn.Linear(16, 4))
|
|
25
|
+
opt = torch.optim.Adam(model.parameters(), lr=0.05)
|
|
26
|
+
for _ in range(300):
|
|
27
|
+
opt.zero_grad()
|
|
28
|
+
torch.nn.functional.cross_entropy(model(x), y).backward()
|
|
29
|
+
opt.step()
|
|
30
|
+
return model, x, y
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
def scripted():
|
|
34
|
+
"""What a good run could look like, so the example works offline."""
|
|
35
|
+
return ScriptedLLM(
|
|
36
|
+
[
|
|
37
|
+
reply("Check behaviour first.", call("forward_accuracy")),
|
|
38
|
+
reply(None, call("describe_tensor", name="2.bias"), call("read_tensor", name="2.bias")),
|
|
39
|
+
reply(
|
|
40
|
+
"Output bias for class 3 is far below the others; set it to their mean.",
|
|
41
|
+
call("checkpoint", label="start"),
|
|
42
|
+
call("python", code="b = W['2.bias']; b[3] = np.delete(b, 3).mean(); b"),
|
|
43
|
+
),
|
|
44
|
+
reply(None, call("forward_accuracy")),
|
|
45
|
+
reply(None, call("submit", answer="Class 3's output bias was -8; reset it to the others' mean.", confidence=0.9)),
|
|
46
|
+
]
|
|
47
|
+
)
|
|
48
|
+
|
|
49
|
+
|
|
50
|
+
def main():
|
|
51
|
+
ap = argparse.ArgumentParser()
|
|
52
|
+
ap.add_argument("--llm", help='e.g. "openai:gpt-5"; omit to use a scripted LLM')
|
|
53
|
+
ap.add_argument("--out", default="runs/fix_quadrants")
|
|
54
|
+
args = ap.parse_args()
|
|
55
|
+
|
|
56
|
+
model, x, y = train()
|
|
57
|
+
with torch.no_grad():
|
|
58
|
+
model[2].bias[3] -= 8.0 # the bug
|
|
59
|
+
|
|
60
|
+
@wsp.probe
|
|
61
|
+
def accuracy():
|
|
62
|
+
"""Per-class accuracy of the network on 4,096 held-out points."""
|
|
63
|
+
with torch.no_grad():
|
|
64
|
+
pred = model(x).argmax(-1)
|
|
65
|
+
return {f"class {c}": round(float((pred[y == c] == c).float().mean()), 3) for c in range(4)}
|
|
66
|
+
|
|
67
|
+
print("before:", accuracy())
|
|
68
|
+
res = wsp.fix(
|
|
69
|
+
model,
|
|
70
|
+
"The classifier never predicts class 3. Find out why and fix it.",
|
|
71
|
+
llm=args.llm or scripted(),
|
|
72
|
+
context=(
|
|
73
|
+
"Network: Linear(2->16) -> ReLU -> Linear(16->4). Input: a point (x0, x1) ~ N(0, I). "
|
|
74
|
+
"Output: logits for the quadrant: class = 2*(x0>0) + (x1>0)."
|
|
75
|
+
),
|
|
76
|
+
probes=[accuracy],
|
|
77
|
+
forward_budget=3,
|
|
78
|
+
verbose=True,
|
|
79
|
+
)
|
|
80
|
+
print("\n" + str(res))
|
|
81
|
+
print("after:", accuracy())
|
|
82
|
+
res.save(args.out)
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
if __name__ == "__main__":
|
|
86
|
+
main()
|
|
@@ -0,0 +1,106 @@
|
|
|
1
|
+
"""Improve a Hugging Face causal LM by editing weights (Experiment 8 style).
|
|
2
|
+
|
|
3
|
+
The LLM gets a Python kernel over the weights and a budget of forward passes:
|
|
4
|
+
one probe generates answers and reports how often they loop, another returns
|
|
5
|
+
only attention patterns (Experiment 7 style).
|
|
6
|
+
|
|
7
|
+
python examples/hf_reduce_loops.py --model Qwen/Qwen2.5-0.5B-Instruct --llm openai:gpt-5 --budget 25
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
import argparse
|
|
11
|
+
|
|
12
|
+
import torch
|
|
13
|
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
14
|
+
|
|
15
|
+
import weightspace as wsp
|
|
16
|
+
|
|
17
|
+
PROMPTS = [
|
|
18
|
+
"What is 17 * 23? Think step by step.",
|
|
19
|
+
"A train travels 60 km in 45 minutes. What is its speed in km/h?",
|
|
20
|
+
"Solve for x: 3x + 7 = 22.",
|
|
21
|
+
"How many prime numbers are there below 30?",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def repetition(ids: list[int], n: int = 8) -> float:
|
|
26
|
+
grams = [tuple(ids[i : i + n]) for i in range(len(ids) - n)]
|
|
27
|
+
return 1 - len(set(grams)) / len(grams) if grams else 0.0
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def main():
|
|
31
|
+
ap = argparse.ArgumentParser()
|
|
32
|
+
ap.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct")
|
|
33
|
+
ap.add_argument("--llm", required=True, help='e.g. "openai:gpt-5", or "dry-run" for a scripted smoke test')
|
|
34
|
+
ap.add_argument("--budget", type=int, default=25)
|
|
35
|
+
ap.add_argument("--max-new-tokens", type=int, default=256)
|
|
36
|
+
ap.add_argument("--out", default="runs/hf_reduce_loops")
|
|
37
|
+
args = ap.parse_args()
|
|
38
|
+
|
|
39
|
+
tok = AutoTokenizer.from_pretrained(args.model)
|
|
40
|
+
model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.float32, attn_implementation="eager").eval()
|
|
41
|
+
|
|
42
|
+
def chat(p):
|
|
43
|
+
if tok.chat_template:
|
|
44
|
+
return tok.apply_chat_template([{"role": "user", "content": p}], tokenize=False, add_generation_prompt=True)
|
|
45
|
+
return p
|
|
46
|
+
|
|
47
|
+
@wsp.probe
|
|
48
|
+
def generate(prompts: list = None, max_new_tokens: int = args.max_new_tokens):
|
|
49
|
+
"""Generate greedy answers for up to 8 prompts (default: a fixed practice set). Returns each answer,
|
|
50
|
+
its 8-gram repetition rate, and whether it hit the token limit."""
|
|
51
|
+
prompts = (prompts or PROMPTS)[:8]
|
|
52
|
+
out = []
|
|
53
|
+
for p in prompts:
|
|
54
|
+
enc = tok(chat(p), return_tensors="pt")
|
|
55
|
+
with torch.no_grad():
|
|
56
|
+
gen = model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
|
|
57
|
+
ids = gen[0, enc.input_ids.shape[1] :].tolist()
|
|
58
|
+
out.append(
|
|
59
|
+
{
|
|
60
|
+
"prompt": p,
|
|
61
|
+
"answer": tok.decode(ids, skip_special_tokens=True)[-1500:],
|
|
62
|
+
"repetition": round(repetition(ids), 3),
|
|
63
|
+
"hit_limit": len(ids) >= max_new_tokens,
|
|
64
|
+
}
|
|
65
|
+
)
|
|
66
|
+
return out
|
|
67
|
+
|
|
68
|
+
@wsp.probe
|
|
69
|
+
def attention(text: str, layer: int):
|
|
70
|
+
"""Attention weights of every head in one layer for a text you write (no outputs are shown)."""
|
|
71
|
+
enc = tok(text, return_tensors="pt")
|
|
72
|
+
with torch.no_grad():
|
|
73
|
+
att = model(**enc, output_attentions=True).attentions[layer][0]
|
|
74
|
+
toks = tok.convert_ids_to_tokens(enc.input_ids[0])
|
|
75
|
+
return {"tokens": toks, "attention[head][query][key]": att.round(decimals=3).tolist()}
|
|
76
|
+
|
|
77
|
+
llm = args.llm
|
|
78
|
+
if llm == "dry-run":
|
|
79
|
+
from weightspace.llm import ScriptedLLM, call, reply
|
|
80
|
+
|
|
81
|
+
llm = ScriptedLLM(
|
|
82
|
+
[
|
|
83
|
+
reply(None, call("list_tensors", pattern="mlp")),
|
|
84
|
+
reply(None, call("forward_generate", prompts=PROMPTS[:1], max_new_tokens=32)),
|
|
85
|
+
reply(None, call("forward_attention", text="2 + 2 = 4", layer=0)),
|
|
86
|
+
reply(None, call("submit", answer="dry run", confidence=0)),
|
|
87
|
+
]
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
res = wsp.improve(
|
|
91
|
+
model,
|
|
92
|
+
"Make the model stop getting stuck in repetition loops on math questions, without hurting its answers.",
|
|
93
|
+
llm=llm,
|
|
94
|
+
context=f"{args.model}: {model.config.num_hidden_layers} layers, hidden size {model.config.hidden_size}. "
|
|
95
|
+
f"Tokenizer vocab {len(tok)}. You may use `tok` (the tokenizer) in Python.",
|
|
96
|
+
probes=[generate, attention],
|
|
97
|
+
forward_budget=args.budget,
|
|
98
|
+
namespace={"tok": tok},
|
|
99
|
+
verbose=True,
|
|
100
|
+
)
|
|
101
|
+
print(res)
|
|
102
|
+
res.save(args.out)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
if __name__ == "__main__":
|
|
106
|
+
main()
|
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
"""Ask an LLM which task a network was trained on, from its weights alone (no tools).
|
|
2
|
+
|
|
3
|
+
Two identical MLPs learn different rules over 8 inputs; the LLM gets one
|
|
4
|
+
network's weights as text and must say which rule it computes.
|
|
5
|
+
|
|
6
|
+
python examples/read_which_task.py --llm openai:gpt-5 --task parity
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import argparse
|
|
10
|
+
|
|
11
|
+
import torch
|
|
12
|
+
|
|
13
|
+
import weightspace as wsp
|
|
14
|
+
|
|
15
|
+
TASKS = {
|
|
16
|
+
"majority": lambda x: (x.sum(1) > 4).long(),
|
|
17
|
+
"first_bit": lambda x: x[:, 0].long(),
|
|
18
|
+
"parity_of_first_two": lambda x: (x[:, 0] != x[:, 1]).long(),
|
|
19
|
+
"last_bit_and_not_third": lambda x: ((x[:, 7] == 1) & (x[:, 2] == 0)).long(),
|
|
20
|
+
}
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def train(rule, seed=0):
|
|
24
|
+
torch.manual_seed(seed)
|
|
25
|
+
x = torch.randint(0, 2, (4096, 8)).float()
|
|
26
|
+
y = rule(x)
|
|
27
|
+
model = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Tanh(), torch.nn.Linear(8, 2))
|
|
28
|
+
opt = torch.optim.Adam(model.parameters(), lr=0.03)
|
|
29
|
+
for _ in range(600):
|
|
30
|
+
opt.zero_grad()
|
|
31
|
+
torch.nn.functional.cross_entropy(model(x), y).backward()
|
|
32
|
+
opt.step()
|
|
33
|
+
return model
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def main():
|
|
37
|
+
ap = argparse.ArgumentParser()
|
|
38
|
+
ap.add_argument("--llm", required=True)
|
|
39
|
+
ap.add_argument("--task", choices=list(TASKS), default="parity_of_first_two")
|
|
40
|
+
args = ap.parse_args()
|
|
41
|
+
|
|
42
|
+
model = train(TASKS[args.task])
|
|
43
|
+
res = wsp.read(
|
|
44
|
+
model,
|
|
45
|
+
"What boolean function of the 8 input bits does this network compute? Answer as a short formula.",
|
|
46
|
+
llm=args.llm,
|
|
47
|
+
context="MLP 8 -> 8 (tanh) -> 2 logits. Inputs are bits in {0,1}. Output 1 means the function is true.",
|
|
48
|
+
text_mode=True,
|
|
49
|
+
precision=3,
|
|
50
|
+
)
|
|
51
|
+
print(res)
|
|
52
|
+
print("truth:", args.task)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
if __name__ == "__main__":
|
|
56
|
+
main()
|
|
@@ -0,0 +1,49 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "weightspace"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Weight-space programming: let an LLM read and write neural network weights directly."
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.10"
|
|
11
|
+
dependencies = ["numpy>=1.24"]
|
|
12
|
+
keywords = ["llm", "interpretability", "model-editing", "weights", "pytorch"]
|
|
13
|
+
classifiers = [
|
|
14
|
+
"Programming Language :: Python :: 3",
|
|
15
|
+
"Intended Audience :: Science/Research",
|
|
16
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
17
|
+
]
|
|
18
|
+
|
|
19
|
+
[project.urls]
|
|
20
|
+
Homepage = "https://github.com/leofmarciano/weightspace"
|
|
21
|
+
Source = "https://github.com/leofmarciano/weightspace"
|
|
22
|
+
|
|
23
|
+
[project.optional-dependencies]
|
|
24
|
+
torch = ["torch>=2.1"]
|
|
25
|
+
hf = ["torch>=2.1", "transformers>=4.40", "safetensors>=0.4"]
|
|
26
|
+
openai = ["openai>=1.40"]
|
|
27
|
+
anthropic = ["anthropic>=0.40"]
|
|
28
|
+
all = ["torch>=2.1", "transformers>=4.40", "safetensors>=0.4", "openai>=1.40", "anthropic>=0.40"]
|
|
29
|
+
dev = ["pytest>=8", "ruff>=0.5"]
|
|
30
|
+
|
|
31
|
+
[project.scripts]
|
|
32
|
+
weightspace = "weightspace.cli:main"
|
|
33
|
+
|
|
34
|
+
[tool.hatch.build.targets.sdist]
|
|
35
|
+
exclude = ["runs", "experiments"]
|
|
36
|
+
|
|
37
|
+
[tool.hatch.build.targets.wheel]
|
|
38
|
+
packages = ["src/weightspace"]
|
|
39
|
+
|
|
40
|
+
[tool.ruff]
|
|
41
|
+
line-length = 130
|
|
42
|
+
target-version = "py310"
|
|
43
|
+
|
|
44
|
+
[tool.ruff.lint]
|
|
45
|
+
select = ["E", "F", "I", "B", "UP"]
|
|
46
|
+
ignore = ["E741", "B023"]
|
|
47
|
+
|
|
48
|
+
[tool.pytest.ini_options]
|
|
49
|
+
testpaths = ["tests"]
|