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.
@@ -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,9 @@
1
+ __pycache__/
2
+ *.egg-info/
3
+ .venv/
4
+ .pytest_cache/
5
+ .ruff_cache/
6
+ dist/
7
+ runs/
8
+ *.npz
9
+ *.safetensors
@@ -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"]