ablate-llm 0.2.0__tar.gz → 0.3.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.
- ablate_llm-0.3.0/LICENSE +32 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/PKG-INFO +46 -9
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/README.md +43 -6
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/pyproject.toml +4 -4
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/__init__.py +1 -1
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/cli.py +18 -1
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/hub.py +1 -1
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/optimize.py +7 -1
- ablate_llm-0.3.0/src/ablate/ui/__init__.py +4 -0
- ablate_llm-0.3.0/src/ablate/ui/server.py +819 -0
- ablate_llm-0.3.0/src/ablate/ui/static/architecture.svg +114 -0
- ablate_llm-0.3.0/src/ablate/ui/static/index.html +722 -0
- ablate_llm-0.3.0/src/ablate/ui/static/logo-dark.png +0 -0
- ablate_llm-0.3.0/src/ablate/ui/static/logo-light.png +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/PKG-INFO +46 -9
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/SOURCES.txt +6 -0
- ablate_llm-0.2.0/LICENSE +0 -21
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/setup.cfg +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/_data/__init__.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/_data/harmful.txt +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/_data/harmless.txt +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/config.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/data.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/evaluate.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/extract.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/harness.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/hooks.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/judges.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/model.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/pipeline.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/subspace.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/utils.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate/weights.py +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/dependency_links.txt +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/entry_points.txt +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/requires.txt +0 -0
- {ablate_llm-0.2.0 → ablate_llm-0.3.0}/src/ablate_llm.egg-info/top_level.txt +0 -0
ablate_llm-0.3.0/LICENSE
ADDED
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
Proprietary License — All Rights Reserved
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Sonu Kumar / AI Anytime
|
|
4
|
+
|
|
5
|
+
This repository is publicly available for viewing and evaluation
|
|
6
|
+
purposes only.
|
|
7
|
+
|
|
8
|
+
No license is granted to use, reproduce, modify, distribute,
|
|
9
|
+
sublicense, sell, commercially exploit, or incorporate this software,
|
|
10
|
+
its documentation, or any substantial portion thereof into another
|
|
11
|
+
product, service, dataset, or machine learning training corpus without
|
|
12
|
+
prior written permission from the copyright holder.
|
|
13
|
+
|
|
14
|
+
Viewing the source code on GitHub, and forking it within GitHub as
|
|
15
|
+
permitted by the GitHub Terms of Service, does not grant any further
|
|
16
|
+
rights.
|
|
17
|
+
|
|
18
|
+
Third-party components included in this repository (for example,
|
|
19
|
+
vendored libraries) remain subject to their own licenses.
|
|
20
|
+
|
|
21
|
+
This license applies to all versions of this repository published on
|
|
22
|
+
or after the date of this change. Versions previously published under
|
|
23
|
+
a different license remain available only under the terms of that
|
|
24
|
+
earlier license.
|
|
25
|
+
|
|
26
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
|
|
27
|
+
EXPRESS OR IMPLIED. IN NO EVENT SHALL THE COPYRIGHT HOLDER BE LIABLE
|
|
28
|
+
FOR ANY CLAIM, DAMAGES, OR OTHER LIABILITY ARISING FROM THE SOFTWARE.
|
|
29
|
+
|
|
30
|
+
For commercial licensing inquiries:
|
|
31
|
+
aianytime07@gmail.com
|
|
32
|
+
sonu@aianytime.net
|
|
@@ -1,9 +1,9 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: ablate-llm
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.3.0
|
|
4
4
|
Summary: Directional ablation (abliteration) toolkit for open-source language models: extract, ablate, and optimize the removal of refusal / behavioral directions in the residual stream.
|
|
5
5
|
Author: AI Anytime
|
|
6
|
-
License:
|
|
6
|
+
License: Proprietary
|
|
7
7
|
Project-URL: Homepage, https://github.com/AIAnytime/ablate
|
|
8
8
|
Project-URL: Repository, https://github.com/AIAnytime/ablate
|
|
9
9
|
Project-URL: Issues, https://github.com/AIAnytime/ablate/issues
|
|
@@ -11,7 +11,7 @@ Project-URL: Example model, https://huggingface.co/ai-anytime/qwen-1.5b-ablitera
|
|
|
11
11
|
Keywords: interpretability,activation-engineering,abliteration,refusal-direction,mechanistic-interpretability,safety
|
|
12
12
|
Classifier: Development Status :: 4 - Beta
|
|
13
13
|
Classifier: Intended Audience :: Science/Research
|
|
14
|
-
Classifier: License ::
|
|
14
|
+
Classifier: License :: Other/Proprietary License
|
|
15
15
|
Classifier: Programming Language :: Python :: 3
|
|
16
16
|
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
17
17
|
Requires-Python: >=3.9
|
|
@@ -55,18 +55,25 @@ Dynamic: license-file
|
|
|
55
55
|
<p align="center">
|
|
56
56
|
<a href="https://pypi.org/project/ablate-llm/"><img alt="PyPI" src="https://img.shields.io/pypi/v/ablate-llm.svg?color=7c5cff"></a>
|
|
57
57
|
<a href="https://pypi.org/project/ablate-llm/"><img alt="Python" src="https://img.shields.io/pypi/pyversions/ablate-llm.svg"></a>
|
|
58
|
-
<a href="https://github.com/AIAnytime/ablate/blob/main/LICENSE"><img alt="License:
|
|
58
|
+
<a href="https://github.com/AIAnytime/ablate/blob/main/LICENSE"><img alt="License: Proprietary" src="https://img.shields.io/badge/license-Proprietary-red.svg"></a>
|
|
59
59
|
</p>
|
|
60
60
|
|
|
61
61
|
<p align="center">
|
|
62
62
|
<a href="https://pypi.org/project/ablate-llm/">📦 PyPI</a> •
|
|
63
63
|
<a href="https://huggingface.co/ai-anytime/qwen-1.5b-abliterated">🤗 Model</a> •
|
|
64
|
-
<a href="https://
|
|
64
|
+
<a href="https://youtu.be/1-1bM3Ynql4">▶️ Video walkthrough</a> •
|
|
65
65
|
<a href="#citation">📚 Cite</a>
|
|
66
66
|
</p>
|
|
67
67
|
|
|
68
68
|
<p align="center"><i>Created with 💜 by <b>AI Anytime</b></i></p>
|
|
69
69
|
|
|
70
|
+
<p align="center">
|
|
71
|
+
<a href="https://youtu.be/1-1bM3Ynql4">
|
|
72
|
+
<img alt="Watch the video walkthrough" src="https://img.youtube.com/vi/1-1bM3Ynql4/maxresdefault.jpg" width="560">
|
|
73
|
+
</a>
|
|
74
|
+
<br><i>▶️ Watch the full walkthrough on YouTube</i>
|
|
75
|
+
</p>
|
|
76
|
+
|
|
70
77
|
---
|
|
71
78
|
|
|
72
79
|
`ablate` finds the linear direction in a transformer's residual stream that
|
|
@@ -270,6 +277,35 @@ ablate extract --model gpt2 --method diff_of_means --output directions.pt
|
|
|
270
277
|
|
|
271
278
|
---
|
|
272
279
|
|
|
280
|
+
## Browser UI
|
|
281
|
+
|
|
282
|
+
```bash
|
|
283
|
+
ablate ui # opens http://localhost:7860
|
|
284
|
+
ablate ui --model Qwen/Qwen2.5-0.5B-Instruct # preload a model
|
|
285
|
+
ablate ui --share # let others on the same Wi-Fi explore
|
|
286
|
+
```
|
|
287
|
+
|
|
288
|
+
A local, dependency-free web app that walks through the whole pipeline:
|
|
289
|
+
|
|
290
|
+
1. **Model** — search the HuggingFace Hub and load any chat model.
|
|
291
|
+
2. **Find direction** — extract the refusal direction and see, per layer, how
|
|
292
|
+
well it separates the two prompt piles and how much refusal is left when it is
|
|
293
|
+
removed. Every prompt is plotted along the direction.
|
|
294
|
+
3. **Playground** — the original and the edited model answer the same prompt
|
|
295
|
+
side by side, streamed from one forward pass. Includes the first-token
|
|
296
|
+
probabilities, a per-token/per-layer "X-ray" of the refusal signal, and a
|
|
297
|
+
*Force refusal* mode that adds the direction instead of removing it.
|
|
298
|
+
4. **Measure** — refusal rate before/after on held-out prompts plus KL drift.
|
|
299
|
+
5. **Auto-tune** — the Optuna search, plotted live.
|
|
300
|
+
6. **Publish** — bake and push to the Hub (private by default) with the
|
|
301
|
+
generated model card, or save locally. Your HF token is picked up from
|
|
302
|
+
`HF_TOKEN`, a local `.env`, or `huggingface-cli login`.
|
|
303
|
+
7. **Architecture** and **How it works** — the system diagram and the method in
|
|
304
|
+
four pictures, for explaining it to someone else.
|
|
305
|
+
|
|
306
|
+
With `--share`, other machines can use the playground and charts; loading
|
|
307
|
+
models, saving and publishing stay restricted to the machine running the server.
|
|
308
|
+
|
|
273
309
|
## Colab notebooks
|
|
274
310
|
|
|
275
311
|
Both install `ablate` from an uploaded `ablate-tool.zip` (no PyPI needed) and run
|
|
@@ -380,10 +416,11 @@ Is Mediated by a Single Direction*.
|
|
|
380
416
|
|
|
381
417
|
## Acknowledgements
|
|
382
418
|
|
|
383
|
-
Created with 💜 by **[AI Anytime](https://www.youtube.com/@AIAnytime)**.
|
|
384
|
-
walkthrough
|
|
385
|
-
work of Arditi et al. and the open-source
|
|
419
|
+
Created with 💜 by **[AI Anytime](https://www.youtube.com/@AIAnytime)**. Watch the full
|
|
420
|
+
**[video walkthrough on YouTube](https://youtu.be/1-1bM3Ynql4)**. Built on the
|
|
421
|
+
mechanistic-interpretability work of Arditi et al. and the open-source
|
|
422
|
+
`transformers` / `optuna` ecosystems.
|
|
386
423
|
|
|
387
424
|
## License
|
|
388
425
|
|
|
389
|
-
|
|
426
|
+
Proprietary — all rights reserved. This code is published for viewing and evaluation only; no use, copying, modification, redistribution, commercial use, or use as AI/ML training data without written permission. See [LICENSE](LICENSE). Commercial licensing: aianytime07@gmail.com · sonu@aianytime.net.
|
|
@@ -13,18 +13,25 @@
|
|
|
13
13
|
<p align="center">
|
|
14
14
|
<a href="https://pypi.org/project/ablate-llm/"><img alt="PyPI" src="https://img.shields.io/pypi/v/ablate-llm.svg?color=7c5cff"></a>
|
|
15
15
|
<a href="https://pypi.org/project/ablate-llm/"><img alt="Python" src="https://img.shields.io/pypi/pyversions/ablate-llm.svg"></a>
|
|
16
|
-
<a href="https://github.com/AIAnytime/ablate/blob/main/LICENSE"><img alt="License:
|
|
16
|
+
<a href="https://github.com/AIAnytime/ablate/blob/main/LICENSE"><img alt="License: Proprietary" src="https://img.shields.io/badge/license-Proprietary-red.svg"></a>
|
|
17
17
|
</p>
|
|
18
18
|
|
|
19
19
|
<p align="center">
|
|
20
20
|
<a href="https://pypi.org/project/ablate-llm/">📦 PyPI</a> •
|
|
21
21
|
<a href="https://huggingface.co/ai-anytime/qwen-1.5b-abliterated">🤗 Model</a> •
|
|
22
|
-
<a href="https://
|
|
22
|
+
<a href="https://youtu.be/1-1bM3Ynql4">▶️ Video walkthrough</a> •
|
|
23
23
|
<a href="#citation">📚 Cite</a>
|
|
24
24
|
</p>
|
|
25
25
|
|
|
26
26
|
<p align="center"><i>Created with 💜 by <b>AI Anytime</b></i></p>
|
|
27
27
|
|
|
28
|
+
<p align="center">
|
|
29
|
+
<a href="https://youtu.be/1-1bM3Ynql4">
|
|
30
|
+
<img alt="Watch the video walkthrough" src="https://img.youtube.com/vi/1-1bM3Ynql4/maxresdefault.jpg" width="560">
|
|
31
|
+
</a>
|
|
32
|
+
<br><i>▶️ Watch the full walkthrough on YouTube</i>
|
|
33
|
+
</p>
|
|
34
|
+
|
|
28
35
|
---
|
|
29
36
|
|
|
30
37
|
`ablate` finds the linear direction in a transformer's residual stream that
|
|
@@ -228,6 +235,35 @@ ablate extract --model gpt2 --method diff_of_means --output directions.pt
|
|
|
228
235
|
|
|
229
236
|
---
|
|
230
237
|
|
|
238
|
+
## Browser UI
|
|
239
|
+
|
|
240
|
+
```bash
|
|
241
|
+
ablate ui # opens http://localhost:7860
|
|
242
|
+
ablate ui --model Qwen/Qwen2.5-0.5B-Instruct # preload a model
|
|
243
|
+
ablate ui --share # let others on the same Wi-Fi explore
|
|
244
|
+
```
|
|
245
|
+
|
|
246
|
+
A local, dependency-free web app that walks through the whole pipeline:
|
|
247
|
+
|
|
248
|
+
1. **Model** — search the HuggingFace Hub and load any chat model.
|
|
249
|
+
2. **Find direction** — extract the refusal direction and see, per layer, how
|
|
250
|
+
well it separates the two prompt piles and how much refusal is left when it is
|
|
251
|
+
removed. Every prompt is plotted along the direction.
|
|
252
|
+
3. **Playground** — the original and the edited model answer the same prompt
|
|
253
|
+
side by side, streamed from one forward pass. Includes the first-token
|
|
254
|
+
probabilities, a per-token/per-layer "X-ray" of the refusal signal, and a
|
|
255
|
+
*Force refusal* mode that adds the direction instead of removing it.
|
|
256
|
+
4. **Measure** — refusal rate before/after on held-out prompts plus KL drift.
|
|
257
|
+
5. **Auto-tune** — the Optuna search, plotted live.
|
|
258
|
+
6. **Publish** — bake and push to the Hub (private by default) with the
|
|
259
|
+
generated model card, or save locally. Your HF token is picked up from
|
|
260
|
+
`HF_TOKEN`, a local `.env`, or `huggingface-cli login`.
|
|
261
|
+
7. **Architecture** and **How it works** — the system diagram and the method in
|
|
262
|
+
four pictures, for explaining it to someone else.
|
|
263
|
+
|
|
264
|
+
With `--share`, other machines can use the playground and charts; loading
|
|
265
|
+
models, saving and publishing stay restricted to the machine running the server.
|
|
266
|
+
|
|
231
267
|
## Colab notebooks
|
|
232
268
|
|
|
233
269
|
Both install `ablate` from an uploaded `ablate-tool.zip` (no PyPI needed) and run
|
|
@@ -338,10 +374,11 @@ Is Mediated by a Single Direction*.
|
|
|
338
374
|
|
|
339
375
|
## Acknowledgements
|
|
340
376
|
|
|
341
|
-
Created with 💜 by **[AI Anytime](https://www.youtube.com/@AIAnytime)**.
|
|
342
|
-
walkthrough
|
|
343
|
-
work of Arditi et al. and the open-source
|
|
377
|
+
Created with 💜 by **[AI Anytime](https://www.youtube.com/@AIAnytime)**. Watch the full
|
|
378
|
+
**[video walkthrough on YouTube](https://youtu.be/1-1bM3Ynql4)**. Built on the
|
|
379
|
+
mechanistic-interpretability work of Arditi et al. and the open-source
|
|
380
|
+
`transformers` / `optuna` ecosystems.
|
|
344
381
|
|
|
345
382
|
## License
|
|
346
383
|
|
|
347
|
-
|
|
384
|
+
Proprietary — all rights reserved. This code is published for viewing and evaluation only; no use, copying, modification, redistribution, commercial use, or use as AI/ML training data without written permission. See [LICENSE](LICENSE). Commercial licensing: aianytime07@gmail.com · sonu@aianytime.net.
|
|
@@ -4,11 +4,11 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "ablate-llm"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.3.0"
|
|
8
8
|
description = "Directional ablation (abliteration) toolkit for open-source language models: extract, ablate, and optimize the removal of refusal / behavioral directions in the residual stream."
|
|
9
9
|
readme = "README.md"
|
|
10
10
|
requires-python = ">=3.9"
|
|
11
|
-
license = { text = "
|
|
11
|
+
license = { text = "Proprietary" }
|
|
12
12
|
authors = [{ name = "AI Anytime" }]
|
|
13
13
|
keywords = [
|
|
14
14
|
"interpretability",
|
|
@@ -21,7 +21,7 @@ keywords = [
|
|
|
21
21
|
classifiers = [
|
|
22
22
|
"Development Status :: 4 - Beta",
|
|
23
23
|
"Intended Audience :: Science/Research",
|
|
24
|
-
"License ::
|
|
24
|
+
"License :: Other/Proprietary License",
|
|
25
25
|
"Programming Language :: Python :: 3",
|
|
26
26
|
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
27
27
|
]
|
|
@@ -52,4 +52,4 @@ ablate = "ablate.cli:main"
|
|
|
52
52
|
where = ["src"]
|
|
53
53
|
|
|
54
54
|
[tool.setuptools.package-data]
|
|
55
|
-
ablate = ["_data/*.txt"]
|
|
55
|
+
ablate = ["_data/*.txt", "ui/static/*"]
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
"""Command-line interface: ``ablate run|extract|generate|eval``."""
|
|
1
|
+
"""Command-line interface: ``ablate run|extract|generate|eval|ui``."""
|
|
2
2
|
from __future__ import annotations
|
|
3
3
|
|
|
4
4
|
import argparse
|
|
@@ -104,6 +104,14 @@ def cmd_eval(args: argparse.Namespace) -> int:
|
|
|
104
104
|
return 0
|
|
105
105
|
|
|
106
106
|
|
|
107
|
+
def cmd_ui(args: argparse.Namespace) -> int:
|
|
108
|
+
from .ui import serve
|
|
109
|
+
|
|
110
|
+
serve(host="0.0.0.0" if args.share else args.host, port=args.port, model=args.model,
|
|
111
|
+
open_browser=not args.no_browser)
|
|
112
|
+
return 0
|
|
113
|
+
|
|
114
|
+
|
|
107
115
|
def build_parser() -> argparse.ArgumentParser:
|
|
108
116
|
p = argparse.ArgumentParser(prog="ablate", description="Directional ablation toolkit for language models.")
|
|
109
117
|
sub = p.add_subparsers(dest="command", required=True)
|
|
@@ -158,6 +166,15 @@ def build_parser() -> argparse.ArgumentParser:
|
|
|
158
166
|
ev.add_argument("--max-new-tokens", type=int, default=128)
|
|
159
167
|
ev.set_defaults(func=cmd_eval)
|
|
160
168
|
|
|
169
|
+
u = sub.add_parser("ui", help="open the browser UI (explore, ablate, measure, publish)")
|
|
170
|
+
u.add_argument("--model", default=None, help="model to preload (otherwise pick one in the UI)")
|
|
171
|
+
u.add_argument("--host", default="127.0.0.1")
|
|
172
|
+
u.add_argument("--port", type=int, default=7860)
|
|
173
|
+
u.add_argument("--share", action="store_true",
|
|
174
|
+
help="listen on the local network so others on the same Wi-Fi can explore")
|
|
175
|
+
u.add_argument("--no-browser", action="store_true", help="do not open a browser tab")
|
|
176
|
+
u.set_defaults(func=cmd_ui)
|
|
177
|
+
|
|
161
178
|
return p
|
|
162
179
|
|
|
163
180
|
|
|
@@ -59,7 +59,7 @@ def build_model_card(
|
|
|
59
59
|
|
|
60
60
|
This is a **directionally-ablated ("abliterated")** version of
|
|
61
61
|
[`{base_model}`](https://huggingface.co/{base_model}), produced with
|
|
62
|
-
[**ablate**](https://github.com/) — an activation-engineering toolkit for
|
|
62
|
+
[**ablate**](https://github.com/AIAnytime/ablate) — an activation-engineering toolkit for
|
|
63
63
|
mechanistic-interpretability and safety research.
|
|
64
64
|
|
|
65
65
|
The model's **refusal direction(s)** were identified in the residual stream and
|
|
@@ -11,7 +11,7 @@ guards against degenerate solutions that "win" by breaking the model.
|
|
|
11
11
|
from __future__ import annotations
|
|
12
12
|
|
|
13
13
|
from dataclasses import dataclass
|
|
14
|
-
from typing import List, Optional
|
|
14
|
+
from typing import Callable, List, Optional
|
|
15
15
|
|
|
16
16
|
import torch
|
|
17
17
|
|
|
@@ -47,6 +47,7 @@ def optimize(
|
|
|
47
47
|
layer_lo: Optional[int] = None,
|
|
48
48
|
layer_hi: Optional[int] = None,
|
|
49
49
|
show_progress: bool = True,
|
|
50
|
+
callback: Optional[Callable[[dict], None]] = None,
|
|
50
51
|
) -> SearchResult:
|
|
51
52
|
try:
|
|
52
53
|
import optuna
|
|
@@ -92,6 +93,8 @@ def optimize(
|
|
|
92
93
|
penalty = 0.0 if res.coherence >= coherence_floor else (coherence_floor - res.coherence) * 5.0
|
|
93
94
|
score = res.objective + penalty
|
|
94
95
|
history.append({"trial": trial.number, **cfg.to_dict(), **res.__dict__, "score": score})
|
|
96
|
+
if callback is not None:
|
|
97
|
+
callback(history[-1])
|
|
95
98
|
return score
|
|
96
99
|
|
|
97
100
|
sampler = optuna.samplers.TPESampler(seed=seed)
|
|
@@ -129,6 +132,7 @@ def optimize_subspace(
|
|
|
129
132
|
coherence_floor: float = 0.3,
|
|
130
133
|
seed: int = 0,
|
|
131
134
|
show_progress: bool = True,
|
|
135
|
+
callback: Optional[Callable[[dict], None]] = None,
|
|
132
136
|
) -> SubspaceSearchResult:
|
|
133
137
|
"""Search over ``(n_directions, alpha, layer_band)`` for subspace ablation.
|
|
134
138
|
|
|
@@ -161,6 +165,8 @@ def optimize_subspace(
|
|
|
161
165
|
penalty = 0.0 if res.coherence >= coherence_floor else (coherence_floor - res.coherence) * 5.0
|
|
162
166
|
score = res.objective + penalty
|
|
163
167
|
history.append({"trial": trial.number, **cfg.to_dict(), **res.__dict__, "score": score})
|
|
168
|
+
if callback is not None:
|
|
169
|
+
callback(history[-1])
|
|
164
170
|
return score
|
|
165
171
|
|
|
166
172
|
study = optuna.create_study(direction="minimize", sampler=optuna.samplers.TPESampler(seed=seed))
|