weightspace 0.1.0__py3-none-any.whl
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/__init__.py +45 -0
- weightspace/adapters.py +112 -0
- weightspace/agent.py +436 -0
- weightspace/cli.py +73 -0
- weightspace/llm/__init__.py +35 -0
- weightspace/llm/anthropic.py +125 -0
- weightspace/llm/base.py +71 -0
- weightspace/llm/mock.py +38 -0
- weightspace/llm/openai.py +108 -0
- weightspace/probes.py +60 -0
- weightspace/prompts.py +43 -0
- weightspace/render.py +151 -0
- weightspace/sandbox.py +68 -0
- weightspace/space.py +319 -0
- weightspace/tasks.py +57 -0
- weightspace-0.1.0.dist-info/METADATA +181 -0
- weightspace-0.1.0.dist-info/RECORD +19 -0
- weightspace-0.1.0.dist-info/WHEEL +4 -0
- weightspace-0.1.0.dist-info/entry_points.txt +2 -0
weightspace/sandbox.py
ADDED
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Persistent Python namespace the LLM uses to analyse and edit weights.
|
|
2
|
+
|
|
3
|
+
This is a convenience for a trusted LLM, not a security boundary: code runs
|
|
4
|
+
in-process. The model object itself is never put in the namespace, so the
|
|
5
|
+
only way to run the network is through the budgeted forward tools.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import ast
|
|
11
|
+
import contextlib
|
|
12
|
+
import io
|
|
13
|
+
import signal
|
|
14
|
+
import threading
|
|
15
|
+
import traceback
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class SandboxTimeout(Exception):
|
|
19
|
+
pass
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class Sandbox:
|
|
23
|
+
def __init__(self, namespace: dict | None = None, *, timeout: float | None = 120, max_output: int = 12000):
|
|
24
|
+
self.ns: dict = {"__name__": "__weightspace__"}
|
|
25
|
+
self.ns.update(namespace or {})
|
|
26
|
+
self.timeout = timeout
|
|
27
|
+
self.max_output = max_output
|
|
28
|
+
|
|
29
|
+
def run(self, code: str) -> str:
|
|
30
|
+
buf = io.StringIO()
|
|
31
|
+
try:
|
|
32
|
+
tree = ast.parse(code, mode="exec")
|
|
33
|
+
last = None
|
|
34
|
+
if tree.body and isinstance(tree.body[-1], ast.Expr):
|
|
35
|
+
last = ast.Expression(tree.body.pop().value)
|
|
36
|
+
with contextlib.redirect_stdout(buf), contextlib.redirect_stderr(buf), self._deadline():
|
|
37
|
+
exec(compile(tree, "<sandbox>", "exec"), self.ns)
|
|
38
|
+
if last is not None:
|
|
39
|
+
value = eval(compile(last, "<sandbox>", "eval"), self.ns)
|
|
40
|
+
if value is not None:
|
|
41
|
+
print(repr(value))
|
|
42
|
+
except SandboxTimeout:
|
|
43
|
+
buf.write(f"\nTimeoutError: code ran longer than {self.timeout}s")
|
|
44
|
+
except BaseException:
|
|
45
|
+
tb = traceback.format_exc().splitlines()
|
|
46
|
+
buf.write("\n".join(l for l in tb if 'File "' not in l or "<sandbox>" in l))
|
|
47
|
+
out = buf.getvalue()
|
|
48
|
+
if len(out) > self.max_output:
|
|
49
|
+
half = self.max_output // 2
|
|
50
|
+
out = out[:half] + f"\n... [{len(out) - self.max_output:,} chars truncated] ...\n" + out[-half:]
|
|
51
|
+
return out or "(no output)"
|
|
52
|
+
|
|
53
|
+
@contextlib.contextmanager
|
|
54
|
+
def _deadline(self):
|
|
55
|
+
if not self.timeout or threading.current_thread() is not threading.main_thread():
|
|
56
|
+
yield
|
|
57
|
+
return
|
|
58
|
+
|
|
59
|
+
def handler(signum, frame):
|
|
60
|
+
raise SandboxTimeout()
|
|
61
|
+
|
|
62
|
+
old = signal.signal(signal.SIGALRM, handler)
|
|
63
|
+
signal.setitimer(signal.ITIMER_REAL, self.timeout)
|
|
64
|
+
try:
|
|
65
|
+
yield
|
|
66
|
+
finally:
|
|
67
|
+
signal.setitimer(signal.ITIMER_REAL, 0)
|
|
68
|
+
signal.signal(signal.SIGALRM, old)
|
weightspace/space.py
ADDED
|
@@ -0,0 +1,319 @@
|
|
|
1
|
+
"""A model-agnostic view of a network's weights as named numpy arrays.
|
|
2
|
+
|
|
3
|
+
The LLM never sees the model object, only tensors. Edits are tracked against
|
|
4
|
+
the original weights so we can report how much changed, roll back, and export
|
|
5
|
+
a sparse patch.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from __future__ import annotations
|
|
9
|
+
|
|
10
|
+
import fnmatch
|
|
11
|
+
from collections.abc import Callable, Iterator, Mapping
|
|
12
|
+
from dataclasses import dataclass, field
|
|
13
|
+
|
|
14
|
+
import numpy as np
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class Store:
|
|
18
|
+
"""Backend holding the live tensors. Subclasses adapt torch, HF, files..."""
|
|
19
|
+
|
|
20
|
+
def names(self) -> list[str]:
|
|
21
|
+
raise NotImplementedError
|
|
22
|
+
|
|
23
|
+
def shape(self, name: str) -> tuple[int, ...]:
|
|
24
|
+
raise NotImplementedError
|
|
25
|
+
|
|
26
|
+
def dtype(self, name: str) -> str:
|
|
27
|
+
raise NotImplementedError
|
|
28
|
+
|
|
29
|
+
def get(self, name: str) -> np.ndarray:
|
|
30
|
+
"""Return a writable float array; edits to it must survive `flush`."""
|
|
31
|
+
raise NotImplementedError
|
|
32
|
+
|
|
33
|
+
def flush(self, names: list[str]) -> None:
|
|
34
|
+
"""Push edits of `names` back to the underlying model, if any."""
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class NumpyStore(Store):
|
|
38
|
+
def __init__(self, tensors: Mapping[str, np.ndarray]):
|
|
39
|
+
self._t = {}
|
|
40
|
+
for k, v in tensors.items():
|
|
41
|
+
arr = np.asarray(v)
|
|
42
|
+
if not np.issubdtype(arr.dtype, np.floating):
|
|
43
|
+
arr = arr.astype(np.float32)
|
|
44
|
+
self._t[k] = arr
|
|
45
|
+
self._original_dtypes = {k: str(np.asarray(v).dtype) for k, v in tensors.items()}
|
|
46
|
+
|
|
47
|
+
def names(self):
|
|
48
|
+
return list(self._t)
|
|
49
|
+
|
|
50
|
+
def shape(self, name):
|
|
51
|
+
return tuple(self._t[name].shape)
|
|
52
|
+
|
|
53
|
+
def dtype(self, name):
|
|
54
|
+
return self._original_dtypes[name]
|
|
55
|
+
|
|
56
|
+
def get(self, name):
|
|
57
|
+
return self._t[name]
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
@dataclass
|
|
61
|
+
class TensorDiff:
|
|
62
|
+
name: str
|
|
63
|
+
numel: int
|
|
64
|
+
changed: int
|
|
65
|
+
l2_delta: float
|
|
66
|
+
max_abs_delta: float
|
|
67
|
+
rel_delta: float
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
@dataclass
|
|
71
|
+
class DiffSummary:
|
|
72
|
+
total_params: int
|
|
73
|
+
changed_params: int
|
|
74
|
+
tensors: list[TensorDiff] = field(default_factory=list)
|
|
75
|
+
|
|
76
|
+
@property
|
|
77
|
+
def fraction(self) -> float:
|
|
78
|
+
return self.changed_params / self.total_params if self.total_params else 0.0
|
|
79
|
+
|
|
80
|
+
def __str__(self) -> str:
|
|
81
|
+
if not self.tensors:
|
|
82
|
+
return "no weights changed"
|
|
83
|
+
lines = [
|
|
84
|
+
f"{self.changed_params:,} of {self.total_params:,} weights changed "
|
|
85
|
+
f"({100 * self.fraction:.4g}%) in {len(self.tensors)} tensor(s)"
|
|
86
|
+
]
|
|
87
|
+
for t in self.tensors:
|
|
88
|
+
lines.append(
|
|
89
|
+
f" {t.name}: {t.changed:,}/{t.numel:,} changed, "
|
|
90
|
+
f"|delta|={t.l2_delta:.4g} ({100 * t.rel_delta:.3g}% of norm), max={t.max_abs_delta:.4g}"
|
|
91
|
+
)
|
|
92
|
+
return "\n".join(lines)
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
class _ReadOnly(Mapping):
|
|
96
|
+
def __init__(self, ws: WeightSpace):
|
|
97
|
+
self._ws = ws
|
|
98
|
+
|
|
99
|
+
def __getitem__(self, name):
|
|
100
|
+
arr = self._ws._store.get(name).view()
|
|
101
|
+
arr.flags.writeable = False
|
|
102
|
+
return arr
|
|
103
|
+
|
|
104
|
+
def __iter__(self):
|
|
105
|
+
return iter(self._ws.names())
|
|
106
|
+
|
|
107
|
+
def __len__(self):
|
|
108
|
+
return len(self._ws.names())
|
|
109
|
+
|
|
110
|
+
def __repr__(self):
|
|
111
|
+
return f"<read-only weights: {len(self)} tensors>"
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
class _Writable(Mapping):
|
|
115
|
+
"""Mapping handed to the sandbox. Touching a tensor snapshots its original."""
|
|
116
|
+
|
|
117
|
+
def __init__(self, ws: WeightSpace):
|
|
118
|
+
self._ws = ws
|
|
119
|
+
|
|
120
|
+
def __getitem__(self, name):
|
|
121
|
+
return self._ws.edit(name)
|
|
122
|
+
|
|
123
|
+
def __setitem__(self, name, value):
|
|
124
|
+
self._ws.set(name, value)
|
|
125
|
+
|
|
126
|
+
def __iter__(self):
|
|
127
|
+
return iter(self._ws.names())
|
|
128
|
+
|
|
129
|
+
def __len__(self):
|
|
130
|
+
return len(self._ws.names())
|
|
131
|
+
|
|
132
|
+
def __repr__(self):
|
|
133
|
+
return f"<writable weights: {len(self)} tensors>"
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
class WeightSpace:
|
|
137
|
+
"""Named float tensors with edit tracking, checkpoints and patches."""
|
|
138
|
+
|
|
139
|
+
def __init__(self, store: Store, *, on_sync: Callable[[], None] | None = None):
|
|
140
|
+
self._store = store
|
|
141
|
+
self._original: dict[str, np.ndarray] = {}
|
|
142
|
+
self._checkpoints: dict[str, dict[str, np.ndarray]] = {}
|
|
143
|
+
self._on_sync = on_sync
|
|
144
|
+
self.R = _ReadOnly(self)
|
|
145
|
+
self.W = _Writable(self)
|
|
146
|
+
|
|
147
|
+
# construction -----------------------------------------------------------------
|
|
148
|
+
@classmethod
|
|
149
|
+
def from_numpy(cls, tensors: Mapping[str, np.ndarray]) -> WeightSpace:
|
|
150
|
+
return cls(NumpyStore(tensors))
|
|
151
|
+
|
|
152
|
+
@classmethod
|
|
153
|
+
def from_torch(cls, module, **kw) -> WeightSpace:
|
|
154
|
+
from .adapters import TorchStore
|
|
155
|
+
|
|
156
|
+
return cls(TorchStore(module, **kw))
|
|
157
|
+
|
|
158
|
+
@classmethod
|
|
159
|
+
def from_hf(cls, model_or_name, **kw) -> WeightSpace:
|
|
160
|
+
from .adapters import load_hf
|
|
161
|
+
|
|
162
|
+
return cls.from_torch(load_hf(model_or_name, **kw))
|
|
163
|
+
|
|
164
|
+
@classmethod
|
|
165
|
+
def from_file(cls, path: str) -> WeightSpace:
|
|
166
|
+
from .adapters import load_file
|
|
167
|
+
|
|
168
|
+
return cls(NumpyStore(load_file(path)))
|
|
169
|
+
|
|
170
|
+
# reading ----------------------------------------------------------------------
|
|
171
|
+
def names(self, pattern: str | None = None) -> list[str]:
|
|
172
|
+
names = self._store.names()
|
|
173
|
+
if pattern:
|
|
174
|
+
names = [n for n in names if fnmatch.fnmatch(n, pattern) or pattern in n]
|
|
175
|
+
return names
|
|
176
|
+
|
|
177
|
+
def shape(self, name: str) -> tuple[int, ...]:
|
|
178
|
+
return self._store.shape(name)
|
|
179
|
+
|
|
180
|
+
def dtype(self, name: str) -> str:
|
|
181
|
+
return self._store.dtype(name)
|
|
182
|
+
|
|
183
|
+
def numel(self, name: str | None = None) -> int:
|
|
184
|
+
if name is not None:
|
|
185
|
+
return int(np.prod(self.shape(name), dtype=np.int64))
|
|
186
|
+
return sum(self.numel(n) for n in self.names())
|
|
187
|
+
|
|
188
|
+
def __getitem__(self, name: str) -> np.ndarray:
|
|
189
|
+
return self.R[name]
|
|
190
|
+
|
|
191
|
+
def __contains__(self, name: str) -> bool:
|
|
192
|
+
return name in self._store.names()
|
|
193
|
+
|
|
194
|
+
def __iter__(self) -> Iterator[str]:
|
|
195
|
+
return iter(self.names())
|
|
196
|
+
|
|
197
|
+
def __len__(self) -> int:
|
|
198
|
+
return len(self.names())
|
|
199
|
+
|
|
200
|
+
def original(self, name: str) -> np.ndarray:
|
|
201
|
+
if name in self._original:
|
|
202
|
+
return self._original[name]
|
|
203
|
+
return self.R[name]
|
|
204
|
+
|
|
205
|
+
# writing ----------------------------------------------------------------------
|
|
206
|
+
def _snapshot(self, name: str) -> None:
|
|
207
|
+
if name not in self._original:
|
|
208
|
+
if name not in self:
|
|
209
|
+
raise KeyError(f"unknown tensor {name!r}")
|
|
210
|
+
self._original[name] = self._store.get(name).copy()
|
|
211
|
+
|
|
212
|
+
def edit(self, name: str) -> np.ndarray:
|
|
213
|
+
"""Return the live, writable array for `name` (in-place edits are tracked)."""
|
|
214
|
+
self._snapshot(name)
|
|
215
|
+
return self._store.get(name)
|
|
216
|
+
|
|
217
|
+
def set(self, name: str, value) -> None:
|
|
218
|
+
arr = self.edit(name)
|
|
219
|
+
value = np.asarray(value, dtype=arr.dtype)
|
|
220
|
+
if value.shape != arr.shape:
|
|
221
|
+
value = np.broadcast_to(value, arr.shape)
|
|
222
|
+
arr[...] = value
|
|
223
|
+
|
|
224
|
+
def zero_(self, pattern: str | None = None) -> None:
|
|
225
|
+
for n in self.names(pattern):
|
|
226
|
+
self.edit(n)[...] = 0
|
|
227
|
+
|
|
228
|
+
def touched(self) -> list[str]:
|
|
229
|
+
return list(self._original)
|
|
230
|
+
|
|
231
|
+
def sync(self) -> None:
|
|
232
|
+
"""Write edits back into the wrapped model (before a forward pass)."""
|
|
233
|
+
self._store.flush(self.touched())
|
|
234
|
+
if self._on_sync:
|
|
235
|
+
self._on_sync()
|
|
236
|
+
|
|
237
|
+
# diff / rollback ----------------------------------------------------------------
|
|
238
|
+
def diff(self) -> DiffSummary:
|
|
239
|
+
out = DiffSummary(total_params=self.numel(), changed_params=0)
|
|
240
|
+
for name, orig in self._original.items():
|
|
241
|
+
cur = self._store.get(name)
|
|
242
|
+
delta = cur.astype(np.float64) - orig.astype(np.float64)
|
|
243
|
+
changed = int(np.count_nonzero(delta))
|
|
244
|
+
if not changed:
|
|
245
|
+
continue
|
|
246
|
+
l2 = float(np.linalg.norm(delta))
|
|
247
|
+
base = float(np.linalg.norm(orig.astype(np.float64)))
|
|
248
|
+
out.tensors.append(
|
|
249
|
+
TensorDiff(
|
|
250
|
+
name=name,
|
|
251
|
+
numel=int(orig.size),
|
|
252
|
+
changed=changed,
|
|
253
|
+
l2_delta=l2,
|
|
254
|
+
max_abs_delta=float(np.abs(delta).max()),
|
|
255
|
+
rel_delta=l2 / base if base else float("inf"),
|
|
256
|
+
)
|
|
257
|
+
)
|
|
258
|
+
out.changed_params += changed
|
|
259
|
+
return out
|
|
260
|
+
|
|
261
|
+
def revert(self, name: str | None = None) -> None:
|
|
262
|
+
names = [name] if name else list(self._original)
|
|
263
|
+
for n in names:
|
|
264
|
+
if n in self._original:
|
|
265
|
+
self._store.get(n)[...] = self._original[n]
|
|
266
|
+
|
|
267
|
+
def checkpoint(self, label: str) -> None:
|
|
268
|
+
self._checkpoints[label] = {n: self._store.get(n).copy() for n in self._original}
|
|
269
|
+
|
|
270
|
+
def restore(self, label: str) -> None:
|
|
271
|
+
if label not in self._checkpoints:
|
|
272
|
+
raise KeyError(f"no checkpoint {label!r}; have {list(self._checkpoints)}")
|
|
273
|
+
saved = self._checkpoints[label]
|
|
274
|
+
for n in self._original:
|
|
275
|
+
self._store.get(n)[...] = saved.get(n, self._original[n])
|
|
276
|
+
|
|
277
|
+
def checkpoints(self) -> list[str]:
|
|
278
|
+
return list(self._checkpoints)
|
|
279
|
+
|
|
280
|
+
def commit(self) -> None:
|
|
281
|
+
"""Accept current weights as the new baseline."""
|
|
282
|
+
self._original.clear()
|
|
283
|
+
self._checkpoints.clear()
|
|
284
|
+
|
|
285
|
+
# patches ----------------------------------------------------------------------
|
|
286
|
+
def patch(self) -> dict[str, np.ndarray]:
|
|
287
|
+
"""Sparse patch: `<name>::idx` flat indices and `<name>::val` new values."""
|
|
288
|
+
out: dict[str, np.ndarray] = {}
|
|
289
|
+
for t in self.diff().tensors:
|
|
290
|
+
cur = self._store.get(t.name).reshape(-1)
|
|
291
|
+
idx = np.flatnonzero(cur != self._original[t.name].reshape(-1))
|
|
292
|
+
out[f"{t.name}::idx"] = idx.astype(np.int64)
|
|
293
|
+
out[f"{t.name}::val"] = cur[idx].copy()
|
|
294
|
+
return out
|
|
295
|
+
|
|
296
|
+
def save_patch(self, path: str) -> None:
|
|
297
|
+
np.savez_compressed(path, **self.patch())
|
|
298
|
+
|
|
299
|
+
def apply_patch(self, patch: str | Mapping[str, np.ndarray]) -> None:
|
|
300
|
+
if isinstance(patch, str):
|
|
301
|
+
with np.load(patch) as f:
|
|
302
|
+
patch = {k: f[k] for k in f.files}
|
|
303
|
+
for key in patch:
|
|
304
|
+
if not key.endswith("::idx"):
|
|
305
|
+
continue
|
|
306
|
+
name = key[: -len("::idx")]
|
|
307
|
+
arr = self.edit(name).reshape(-1)
|
|
308
|
+
arr[patch[key]] = patch[f"{name}::val"]
|
|
309
|
+
|
|
310
|
+
def to_numpy(self) -> dict[str, np.ndarray]:
|
|
311
|
+
return {n: self._store.get(n).copy() for n in self.names()}
|
|
312
|
+
|
|
313
|
+
def save(self, path: str) -> None:
|
|
314
|
+
from .adapters import save_file
|
|
315
|
+
|
|
316
|
+
save_file(self.to_numpy(), path)
|
|
317
|
+
|
|
318
|
+
def __repr__(self) -> str:
|
|
319
|
+
return f"WeightSpace({len(self)} tensors, {self.numel():,} weights)"
|
weightspace/tasks.py
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
1
|
+
"""One-call entry points: read / fix / write / improve."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from collections.abc import Mapping
|
|
6
|
+
|
|
7
|
+
from .agent import Result, Session
|
|
8
|
+
from .space import WeightSpace
|
|
9
|
+
|
|
10
|
+
|
|
11
|
+
def as_weightspace(model) -> WeightSpace:
|
|
12
|
+
"""Wrap a WeightSpace, torch module, HF model, dict of arrays, or weights file path."""
|
|
13
|
+
if isinstance(model, WeightSpace):
|
|
14
|
+
return model
|
|
15
|
+
if isinstance(model, str):
|
|
16
|
+
return WeightSpace.from_file(model)
|
|
17
|
+
if isinstance(model, Mapping):
|
|
18
|
+
return WeightSpace.from_numpy(model)
|
|
19
|
+
try:
|
|
20
|
+
import torch
|
|
21
|
+
|
|
22
|
+
if isinstance(model, torch.nn.Module):
|
|
23
|
+
return WeightSpace.from_torch(model)
|
|
24
|
+
except ImportError:
|
|
25
|
+
pass
|
|
26
|
+
raise TypeError(f"cannot build a WeightSpace from {type(model).__name__}")
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _run(mode, model, task, llm, text_mode=False, **kw) -> Result:
|
|
30
|
+
ws = as_weightspace(model)
|
|
31
|
+
text_kw = {k: kw.pop(k) for k in ("precision", "max_chars") if k in kw}
|
|
32
|
+
s = Session(ws, llm, mode=mode, **kw)
|
|
33
|
+
return s.run_text(task, **text_kw) if text_mode else s.run(task)
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def read(model, question: str, llm, **kw) -> Result:
|
|
37
|
+
"""Answer a question about a network from its weights. Never edits."""
|
|
38
|
+
return _run("read", model, question, llm, **kw)
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def fix(model, problem: str, llm, **kw) -> Result:
|
|
42
|
+
"""Find and fix the cause of a described failure by editing weights."""
|
|
43
|
+
return _run("fix", model, problem, llm, **kw)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def write(model, spec: str, llm, *, zero_init: bool = True, **kw) -> Result:
|
|
47
|
+
"""Program a network from scratch (weights zeroed first unless zero_init=False)."""
|
|
48
|
+
ws = as_weightspace(model)
|
|
49
|
+
if zero_init:
|
|
50
|
+
ws.zero_()
|
|
51
|
+
ws.commit()
|
|
52
|
+
return _run("write", ws, spec, llm, **kw)
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def improve(model, goal: str, llm, **kw) -> Result:
|
|
56
|
+
"""Make a network better at a goal, usually with probes and a forward budget."""
|
|
57
|
+
return _run("improve", model, goal, llm, **kw)
|
|
@@ -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,19 @@
|
|
|
1
|
+
weightspace/__init__.py,sha256=mq80-2uZ7EZ6PwREbmNdtdjnaxiwFhm1Gp-QkEoJjyw,1168
|
|
2
|
+
weightspace/adapters.py,sha256=lAv6KGTDso45GtdExduAwq1GkzWlVtwnbjGCC4iO2vk,3742
|
|
3
|
+
weightspace/agent.py,sha256=BMAhUACTnmu3MnhUdVfjrcJrEDUmB9v7dKDHszkB864,18361
|
|
4
|
+
weightspace/cli.py,sha256=RVOdV7j_iWvZKU23BXnGIl5E-Ejf3SnrYvJPAJazpOw,2646
|
|
5
|
+
weightspace/probes.py,sha256=vXGFY_qMEGQDDXnbFPV1TUo3GToScaXZXbuB4hmFV3U,2116
|
|
6
|
+
weightspace/prompts.py,sha256=4Q5RNzZnM0Kg9jajfIMDPC9mjllU1hpxIGuM4FwjpuE,2523
|
|
7
|
+
weightspace/render.py,sha256=5-Dv2I4pIidqZ-DloY6fejTEmpaklIMaFEN4MXvFle8,5419
|
|
8
|
+
weightspace/sandbox.py,sha256=GjBoU1PANP37PSe15UhtgIjK9bigUmGFRyTTNtJpEzU,2452
|
|
9
|
+
weightspace/space.py,sha256=foJz4re5U7LDme155Y-cnr3PasNAIkpfwxJgcA-byPg,10404
|
|
10
|
+
weightspace/tasks.py,sha256=OBoe4plzggy1q4Gtyh0MaEB_N0kzjeWEEp-fAPmjbc8,1932
|
|
11
|
+
weightspace/llm/__init__.py,sha256=iezqhDRZGXJfTiyTj9KFUj2U7BjXXQ9A3Oqiavbafdk,1184
|
|
12
|
+
weightspace/llm/anthropic.py,sha256=a4piOrnNmtYyWBMwt6zZFgA0tT8SWDiBv61mB2FCMBk,5311
|
|
13
|
+
weightspace/llm/base.py,sha256=3wguEIsIbx6e5dzX3yVgM_LJYja7ePBEeZydAxPLH8E,2118
|
|
14
|
+
weightspace/llm/mock.py,sha256=p7qva2VmI1Kz5X6rW-VItMjfGJu47IszYK0f5cHGFRM,1188
|
|
15
|
+
weightspace/llm/openai.py,sha256=NhJXV57430F3Y7EMWr6W5Iodzk9Rg0YjwXYBQtlxedQ,4063
|
|
16
|
+
weightspace-0.1.0.dist-info/METADATA,sha256=_Kcg41CQQYgIkEydrLef5C6wXeyLNrGEXtqNWNLmohI,7712
|
|
17
|
+
weightspace-0.1.0.dist-info/WHEEL,sha256=2iZrJkDIkeDh0k9SikKMqFuRP-ua6mb4Ir1AsER4BEw,87
|
|
18
|
+
weightspace-0.1.0.dist-info/entry_points.txt,sha256=UrI1f3liF2AwZmh_RRpv3-NAEdb0AqzgZBFzmeG5cNE,53
|
|
19
|
+
weightspace-0.1.0.dist-info/RECORD,,
|