oumi-ai/weights-lab

An agent skill and Python worker that let your coding agent experiment on an open
language model: load it once, edit its weights reversibly, measure the effect, and
export the result. It is inspired by
weight-space programming
(Takehi, Teoh, Zhang, and Khattab), in which a general-purpose LLM reads and writes a
network's weights directly instead of relying on an optimizer such as gradient
descent. Every snippet's source, output, and forward-pass count is saved.
Version 0.1 is experimental, and snippets run as trusted Python with your filesystem
and network access.
Tested on CPU, Apple MPS, NVIDIA H200, and SkyPilot with Slurm.
A deliberately silly weight edit, with exact undo. Try the example.
## Install in your coding agent
Requires Linux or macOS, git, and uv.
### Claude Code
sh
claude plugin marketplace add oumi-ai/weights-lab
claude plugin install oumi-weights-lab@oumi-weights-lab
The plugin includes the oumi-weights-lab CLI. Its first run downloads PyTorch, so
allow a few minutes.
### Codex
sh
codex plugin marketplace add oumi-ai/weights-lab
codex plugin add oumi-weights-lab@oumi-weights-lab
uv tool install git+https://github.com/oumi-ai/weights-lab
### Cursor, GitHub Copilot CLI, OpenCode, and other skill-compatible agents
Copy the whole skill folder, then install the CLI:
sh
git clone https://github.com/oumi-ai/weights-lab
mkdir -p ~/.agents/skills
cp -R weights-lab/skills/oumi-weights-lab ~/.agents/skills/
uv tool install git+https://github.com/oumi-ai/weights-lab
Then ask your agent, for example: "Use oumi-weights-lab to check whether zeroing
layer 3's MLP in Qwen/Qwen3-0.6B changes its answers to the questions in caps.jsonl."
To load LoRA adapters, or Hugging Face datasets with lab.hfrows, install the CLI
as uv tool install 'oumi-weights-lab[adapters,datasets] @ git+https://github.com/oumi-ai/weights-lab'.
The Claude Code plugin includes both.
Append @ to a git+https:// URL to pin a release.
## Use the CLI directly
sh
oumi-weights-lab start --session sessions/first --model HuggingFaceTB/SmolLM2-135M --device cpu
oumi-weights-lab send --session sessions/first --file edit.py --follow
oumi-weights-lab stop --session sessions/first
oumi-weights-lab backup --session sessions/first --output sessions/first.tar.gz
With edit.py:
python
rows = [{"prompt": "France's capital is", "options": [" Paris", " Rome"], "answer": 0}]
baseline = lab.evaluate(rows, name="baseline")
name = lab.parts("mlpweight")[0]["name"]
lab.zero(name)
edited = lab.evaluate(rows, name="zeroed")
print(lab.diff(baseline, edited))
lab.undo()
checkpoint = lab.export("restored") # saved under artifacts/
Snippets share lab, model, tokenizer, Path, and artifacts, and their state