strands-slm 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,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 cagataycali
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,164 @@
1
+ Metadata-Version: 2.4
2
+ Name: strands-slm
3
+ Version: 0.1.0
4
+ Summary: Self-learning model: a Strands-Agents-expert Qwen3-VL-2B whose weights keep changing at inference — surprise-gated, bounded, with a provable off-switch.
5
+ Author: cagataycali
6
+ License: MIT
7
+ Project-URL: Homepage, https://github.com/cagataycali/slm
8
+ Project-URL: Repository, https://github.com/cagataycali/slm
9
+ Keywords: test-time-learning,continual-learning,fast-weights,lora,qwen,strands-agents,self-learning
10
+ Classifier: Programming Language :: Python :: 3
11
+ Classifier: License :: OSI Approved :: MIT License
12
+ Classifier: Intended Audience :: Science/Research
13
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
14
+ Requires-Python: >=3.9
15
+ Description-Content-Type: text/markdown
16
+ License-File: LICENSE
17
+ Requires-Dist: torch>=2.1
18
+ Requires-Dist: transformers>=4.45
19
+ Requires-Dist: huggingface_hub>=0.20
20
+ Provides-Extra: train
21
+ Requires-Dist: peft>=0.10; extra == "train"
22
+ Requires-Dist: accelerate; extra == "train"
23
+ Provides-Extra: dev
24
+ Requires-Dist: pytest; extra == "dev"
25
+ Dynamic: license-file
26
+
27
+ # slm
28
+
29
+ **A model whose weights change while it runs.**
30
+ Predict, get surprised, rewrite a small bounded part of yourself, never forget the base.
31
+
32
+ Every LLM you have used is frozen at deployment. `slm` wraps a frozen
33
+ [Qwen3-VL-2B post-tuned on the strands-agents codebase](https://huggingface.co/cagataydev/strands-qwen3-vl-2b)
34
+ with a plastic layer that keeps learning at inference — with a provable off-switch.
35
+
36
+ ```bash
37
+ pip install strands-slm
38
+ ```
39
+
40
+ ## Quickstart
41
+
42
+ As a [Strands Agents](https://github.com/strands-agents) model provider — every turn can change the weights:
43
+
44
+ ```python
45
+ from strands import Agent
46
+ from strands_tools import shell
47
+ from slm import SLM
48
+
49
+ model = SLM("cagataydev/strands-qwen3-vl-2b")
50
+ agent = Agent(tools=[shell], model=model)
51
+
52
+ agent("use the shell tool to run: echo hello") # this turn updated the weights
53
+ ```
54
+
55
+ Or drive the learning loop directly:
56
+
57
+ ```python
58
+ from slm import StrandsPlasticQwen
59
+
60
+ m = StrandsPlasticQwen.from_pretrained()
61
+ print(m.chat("How do I create a custom tool in Strands Agents?"))
62
+
63
+ for doc in your_stream:
64
+ m.observe(doc, learn=True) # predicts; if surprised, rewrites its fast weights
65
+ m.reset() # bit-exact back to the base
66
+ ```
67
+
68
+ **See it happen: [demo.ipynb](demo.ipynb)** — ask the model a question it cannot
69
+ know, let it read documents (pure inference), ask again — it knows. Then reset,
70
+ and it forgets. Executed outputs and plots embedded; validated on an L40S:
71
+ P(correct) 0.09 → 0.74, greedy answers 3/3, reset Δlogits = 0.
72
+
73
+ ## How it works
74
+
75
+ ```
76
+ frozen Qwen3-VL-2B instinct — never updated, cannot forget
77
+ + strands LoRA (merged) slow: post-tuned strands-agents expertise
78
+ + plastic LoRA fast: ~1.6M params over the frozen 2.13B,
79
+ updated on every observation at inference
80
+ + surprise gate learn only when prediction error spikes
81
+ + EMA decay bounded plasticity — learns AND retains
82
+
83
+ loss = next-observation prediction error (the free label from reality)
84
+ ```
85
+
86
+ ## Results
87
+
88
+ Measured on a single GPU, seed-replicated. The base model is never updated.
89
+
90
+ | claim | evidence |
91
+ |---|---|
92
+ | Domain expert | strands probe NLL 4.85 → 2.22, 8/8 probes improved |
93
+ | Learns while running | continual OOD stream NLL 6.18 → 5.37, pure inference |
94
+ | Does not forget | strands expertise after OOD learning: Δ −0.01 |
95
+ | Agent competence grows | held-out tasks 0/4 → 4/4 after 18 curated lessons, 5/5 seeds |
96
+ | Fact memory | 15/15 facts at 100% verbatim recall |
97
+ | Fleet learning | two agents' experience files summed losslessly |
98
+ | Persistence | experience survives process death bit-exact |
99
+ | Provable off-switch | `reset()` is bit-identical to the base, Δlogits = 0 |
100
+ | Cost | +0.11–0.25 s/turn learning overhead |
101
+
102
+ The stability–plasticity dial, measured (OOD baseline NLL 4.23):
103
+
104
+ | lr | EMA decay | OOD gain | retention Δ | |
105
+ |---|---|---|---|---|
106
+ | 2e-3 | 0.98 | +0.05 | +0.00 | too timid |
107
+ | 8e-3 | 0.98 | +0.89 | +0.03 | the sweet spot |
108
+ | 1e-2 | none | +3.30 | +7.09 | forgets the base |
109
+
110
+ ## API
111
+
112
+ | method | what it does |
113
+ |---|---|
114
+ | `SLM(model_id, plasticity="high", placement="deep")` | Strands provider; agent turns learn automatically |
115
+ | `.teach(prompt, response)` | curated lesson: bind a future query to a desired response |
116
+ | `.observe(text, learn=True)` | free-form learning; returns pre-update surprise (NLL) |
117
+ | `.consolidate(epochs=5)` | sleep phase: replay the lesson buffer, harden weak memories |
118
+ | `.revise(prompt, old, new)` | targeted unlearning: flip a consolidated belief |
119
+ | `.save_fast_weights(path)` / `.load_fast_weights(path)` | persist or restore acquired experience |
120
+ | `.merge_experience(paths)` | fleet learning: compose multiple agents' experience files |
121
+ | `.reset()` | the off-switch — exactly the base model again |
122
+ | `.surprise_log` | (turn, NLL) history — watch it learn |
123
+
124
+ ## What we learned building it
125
+
126
+ 1. Placement determines what can be learned: attention q/v LoRA stores bindings
127
+ about 4x more sample-efficiently than the LM head.
128
+ 2. There is a free-learning regime (deep placement, lr 2e-2, decay 0.999):
129
+ skill acquisition at zero retention cost.
130
+ 3. You retrieve in the format you learned — render lessons through the real
131
+ chat template or the knowledge is invisible at inference.
132
+ 4. Curation is the difference between experience and learning: raw feedback
133
+ transcripts teach nothing; distilled (task → corrected response) pairs
134
+ take held-out competence from 0/4 to 4/4.
135
+ 5. Interleave or lose it: sequential lessons evict each other; replay makes
136
+ them coexist. Sleep-style consolidation hardens weak memories.
137
+ 6. Belief revision is a terminal operation: whatever is learned last in a
138
+ semantic neighborhood wins — order lessons before the revision.
139
+
140
+ ## Honest limitations
141
+
142
+ - A bolt-on linear memory degrades single-prompt in-context recall — softmax
143
+ attention is already the better mechanism there. The win is persistent
144
+ cross-sequence adaptation, which the context window cannot retain.
145
+ - About a third of naive test-time-training gains in the literature are pure
146
+ calibration (even zero-information targets help an over-confident head).
147
+ Our evals control for this with an information-ladder baseline.
148
+ - Composition (learned schema x unseen entity) plateaus near 67% at 2B.
149
+ - All findings are at 2B scale; scaling behavior is unknown.
150
+
151
+ ## Reproduce the post-tune
152
+
153
+ ```bash
154
+ pip install "strands-slm[train]"
155
+ python scripts/build_corpus.py # strands-agents repos -> corpus.jsonl
156
+ python scripts/train_lora.py --steps 1200 --bs 2 --accum 4 --lr 1e-4
157
+ python scripts/eval_strands.py # base vs tuned probes
158
+ ```
159
+
160
+ Private HF repos need `HF_TOKEN` in the environment, or pass `token=`.
161
+
162
+ ## License
163
+
164
+ MIT
@@ -0,0 +1,138 @@
1
+ # slm
2
+
3
+ **A model whose weights change while it runs.**
4
+ Predict, get surprised, rewrite a small bounded part of yourself, never forget the base.
5
+
6
+ Every LLM you have used is frozen at deployment. `slm` wraps a frozen
7
+ [Qwen3-VL-2B post-tuned on the strands-agents codebase](https://huggingface.co/cagataydev/strands-qwen3-vl-2b)
8
+ with a plastic layer that keeps learning at inference — with a provable off-switch.
9
+
10
+ ```bash
11
+ pip install strands-slm
12
+ ```
13
+
14
+ ## Quickstart
15
+
16
+ As a [Strands Agents](https://github.com/strands-agents) model provider — every turn can change the weights:
17
+
18
+ ```python
19
+ from strands import Agent
20
+ from strands_tools import shell
21
+ from slm import SLM
22
+
23
+ model = SLM("cagataydev/strands-qwen3-vl-2b")
24
+ agent = Agent(tools=[shell], model=model)
25
+
26
+ agent("use the shell tool to run: echo hello") # this turn updated the weights
27
+ ```
28
+
29
+ Or drive the learning loop directly:
30
+
31
+ ```python
32
+ from slm import StrandsPlasticQwen
33
+
34
+ m = StrandsPlasticQwen.from_pretrained()
35
+ print(m.chat("How do I create a custom tool in Strands Agents?"))
36
+
37
+ for doc in your_stream:
38
+ m.observe(doc, learn=True) # predicts; if surprised, rewrites its fast weights
39
+ m.reset() # bit-exact back to the base
40
+ ```
41
+
42
+ **See it happen: [demo.ipynb](demo.ipynb)** — ask the model a question it cannot
43
+ know, let it read documents (pure inference), ask again — it knows. Then reset,
44
+ and it forgets. Executed outputs and plots embedded; validated on an L40S:
45
+ P(correct) 0.09 → 0.74, greedy answers 3/3, reset Δlogits = 0.
46
+
47
+ ## How it works
48
+
49
+ ```
50
+ frozen Qwen3-VL-2B instinct — never updated, cannot forget
51
+ + strands LoRA (merged) slow: post-tuned strands-agents expertise
52
+ + plastic LoRA fast: ~1.6M params over the frozen 2.13B,
53
+ updated on every observation at inference
54
+ + surprise gate learn only when prediction error spikes
55
+ + EMA decay bounded plasticity — learns AND retains
56
+
57
+ loss = next-observation prediction error (the free label from reality)
58
+ ```
59
+
60
+ ## Results
61
+
62
+ Measured on a single GPU, seed-replicated. The base model is never updated.
63
+
64
+ | claim | evidence |
65
+ |---|---|
66
+ | Domain expert | strands probe NLL 4.85 → 2.22, 8/8 probes improved |
67
+ | Learns while running | continual OOD stream NLL 6.18 → 5.37, pure inference |
68
+ | Does not forget | strands expertise after OOD learning: Δ −0.01 |
69
+ | Agent competence grows | held-out tasks 0/4 → 4/4 after 18 curated lessons, 5/5 seeds |
70
+ | Fact memory | 15/15 facts at 100% verbatim recall |
71
+ | Fleet learning | two agents' experience files summed losslessly |
72
+ | Persistence | experience survives process death bit-exact |
73
+ | Provable off-switch | `reset()` is bit-identical to the base, Δlogits = 0 |
74
+ | Cost | +0.11–0.25 s/turn learning overhead |
75
+
76
+ The stability–plasticity dial, measured (OOD baseline NLL 4.23):
77
+
78
+ | lr | EMA decay | OOD gain | retention Δ | |
79
+ |---|---|---|---|---|
80
+ | 2e-3 | 0.98 | +0.05 | +0.00 | too timid |
81
+ | 8e-3 | 0.98 | +0.89 | +0.03 | the sweet spot |
82
+ | 1e-2 | none | +3.30 | +7.09 | forgets the base |
83
+
84
+ ## API
85
+
86
+ | method | what it does |
87
+ |---|---|
88
+ | `SLM(model_id, plasticity="high", placement="deep")` | Strands provider; agent turns learn automatically |
89
+ | `.teach(prompt, response)` | curated lesson: bind a future query to a desired response |
90
+ | `.observe(text, learn=True)` | free-form learning; returns pre-update surprise (NLL) |
91
+ | `.consolidate(epochs=5)` | sleep phase: replay the lesson buffer, harden weak memories |
92
+ | `.revise(prompt, old, new)` | targeted unlearning: flip a consolidated belief |
93
+ | `.save_fast_weights(path)` / `.load_fast_weights(path)` | persist or restore acquired experience |
94
+ | `.merge_experience(paths)` | fleet learning: compose multiple agents' experience files |
95
+ | `.reset()` | the off-switch — exactly the base model again |
96
+ | `.surprise_log` | (turn, NLL) history — watch it learn |
97
+
98
+ ## What we learned building it
99
+
100
+ 1. Placement determines what can be learned: attention q/v LoRA stores bindings
101
+ about 4x more sample-efficiently than the LM head.
102
+ 2. There is a free-learning regime (deep placement, lr 2e-2, decay 0.999):
103
+ skill acquisition at zero retention cost.
104
+ 3. You retrieve in the format you learned — render lessons through the real
105
+ chat template or the knowledge is invisible at inference.
106
+ 4. Curation is the difference between experience and learning: raw feedback
107
+ transcripts teach nothing; distilled (task → corrected response) pairs
108
+ take held-out competence from 0/4 to 4/4.
109
+ 5. Interleave or lose it: sequential lessons evict each other; replay makes
110
+ them coexist. Sleep-style consolidation hardens weak memories.
111
+ 6. Belief revision is a terminal operation: whatever is learned last in a
112
+ semantic neighborhood wins — order lessons before the revision.
113
+
114
+ ## Honest limitations
115
+
116
+ - A bolt-on linear memory degrades single-prompt in-context recall — softmax
117
+ attention is already the better mechanism there. The win is persistent
118
+ cross-sequence adaptation, which the context window cannot retain.
119
+ - About a third of naive test-time-training gains in the literature are pure
120
+ calibration (even zero-information targets help an over-confident head).
121
+ Our evals control for this with an information-ladder baseline.
122
+ - Composition (learned schema x unseen entity) plateaus near 67% at 2B.
123
+ - All findings are at 2B scale; scaling behavior is unknown.
124
+
125
+ ## Reproduce the post-tune
126
+
127
+ ```bash
128
+ pip install "strands-slm[train]"
129
+ python scripts/build_corpus.py # strands-agents repos -> corpus.jsonl
130
+ python scripts/train_lora.py --steps 1200 --bs 2 --accum 4 --lr 1e-4
131
+ python scripts/eval_strands.py # base vs tuned probes
132
+ ```
133
+
134
+ Private HF repos need `HF_TOKEN` in the environment, or pass `token=`.
135
+
136
+ ## License
137
+
138
+ MIT
@@ -0,0 +1,35 @@
1
+ [build-system]
2
+ requires = ["setuptools>=61"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "strands-slm"
7
+ version = "0.1.0"
8
+ description = "Self-learning model: a Strands-Agents-expert Qwen3-VL-2B whose weights keep changing at inference — surprise-gated, bounded, with a provable off-switch."
9
+ readme = "README.md"
10
+ requires-python = ">=3.9"
11
+ license = {text = "MIT"}
12
+ authors = [{name = "cagataycali"}]
13
+ keywords = ["test-time-learning", "continual-learning", "fast-weights", "lora", "qwen", "strands-agents", "self-learning"]
14
+ classifiers = [
15
+ "Programming Language :: Python :: 3",
16
+ "License :: OSI Approved :: MIT License",
17
+ "Intended Audience :: Science/Research",
18
+ "Topic :: Scientific/Engineering :: Artificial Intelligence",
19
+ ]
20
+ dependencies = [
21
+ "torch>=2.1",
22
+ "transformers>=4.45",
23
+ "huggingface_hub>=0.20",
24
+ ]
25
+
26
+ [project.optional-dependencies]
27
+ train = ["peft>=0.10", "accelerate"]
28
+ dev = ["pytest"]
29
+
30
+ [project.urls]
31
+ Homepage = "https://github.com/cagataycali/slm"
32
+ Repository = "https://github.com/cagataycali/slm"
33
+
34
+ [tool.setuptools.packages.find]
35
+ include = ["slm*"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,32 @@
1
+ """
2
+ slm — self-learning model.
3
+
4
+ A Strands-Agents-expert Qwen3-VL-2B that keeps learning after deployment:
5
+
6
+ frozen Qwen3-VL-2B (instinct — never updated, can't forget)
7
+ + strands LoRA (merged) (SLOW: post-tuned strands-agents expertise)
8
+ + plastic LoRA on lm_head (FAST: surprise-gated, EMA-decayed,
9
+ updated at inference — with a provable off-switch)
10
+
11
+ Quick start:
12
+ from slm import StrandsPlasticQwen
13
+
14
+ m = StrandsPlasticQwen.from_pretrained() # cagataydev/strands-qwen3-vl-2b
15
+ print(m.chat("How do I create a custom tool in Strands Agents?"))
16
+
17
+ for doc in your_stream:
18
+ m.observe(doc, learn=True) # predicts; if surprised, rewrites fast weights
19
+ m.reset() # exactly back to the strands-expert base
20
+ """
21
+ from .qwen import StrandsPlasticQwen, DEFAULT_MODEL
22
+
23
+
24
+ def __getattr__(name):
25
+ # Lazy import: SLM needs strands-agents installed
26
+ if name == "SLM":
27
+ from .strands_model import SLM
28
+ return SLM
29
+ raise AttributeError(f"module 'slm' has no attribute {name!r}")
30
+
31
+ __version__ = "0.1.0"
32
+ __all__ = ["StrandsPlasticQwen", "SLM", "DEFAULT_MODEL", "__version__"]
@@ -0,0 +1,158 @@
1
+ """
2
+ slm.qwen — self-learning Qwen3-VL runtimes (optional extra: pip install self-learning-model[qwen]).
3
+
4
+ Two runtimes over a frozen (or merged strands-expert) Qwen3-VL-2B:
5
+
6
+ StrandsPlasticQwen — the Strands-Agents-expert Qwen (post-tuned, merged) with a
7
+ plastic LoRA head on lm_head that keeps learning at inference:
8
+ surprise-gated SGD + EMA decay (bounded plasticity), provable off-switch.
9
+
10
+ Validated (see README.md / PROOF.md in the repo):
11
+ * continual OOD stream: NLL drops online while base knowledge is retained
12
+ * reset() restores the base exactly (Δlogits = 0)
13
+
14
+ Usage:
15
+ from slm.qwen import StrandsPlasticQwen
16
+ m = StrandsPlasticQwen.from_pretrained() # default: strands-expert model
17
+ print(m.chat("How do I create a custom tool in Strands Agents?"))
18
+ m.observe(new_docs, learn=True) # self-learn after deployment
19
+ m.reset() # off-switch
20
+
21
+ Requires: torch, transformers (installed via the [qwen] extra). Private HF repos
22
+ need HF_TOKEN in the environment.
23
+ """
24
+ import os
25
+
26
+ DEFAULT_MODEL = "cagataydev/strands-qwen3-vl-2b"
27
+
28
+
29
+ def _require_torch():
30
+ try:
31
+ import torch # noqa
32
+ import transformers # noqa
33
+ except ImportError as e:
34
+ raise ImportError(
35
+ "slm.qwen needs the optional deps: pip install 'self-learning-model[qwen]'"
36
+ ) from e
37
+
38
+
39
+ class StrandsPlasticQwen:
40
+ """Strands-expert Qwen3-VL-2B + fast plastic LoRA head (self-learning at inference)."""
41
+
42
+ def __init__(self, model, tok, head, lr=8e-3, decay=0.98, k_gate=0.0):
43
+ import torch
44
+ self.model, self.tok, self.head = model, tok, head
45
+ self.opt = torch.optim.SGD([head.A, head.B], lr=lr)
46
+ self.decay, self.k_gate = decay, k_gate
47
+ self.mean, self.beta = None, 0.9
48
+ self.device = next(model.parameters()).device
49
+
50
+ # ---------------- constructors ----------------
51
+ @classmethod
52
+ def from_pretrained(cls, model_id=DEFAULT_MODEL, device="cuda", r_fast=16,
53
+ token=None, **kw):
54
+ """Load the merged strands-expert model (or any Qwen3-VL id) + attach
55
+ the fast plastic head. Private repos: pass token= or set HF_TOKEN."""
56
+ _require_torch()
57
+ import torch
58
+ import torch.nn as nn
59
+ from transformers import AutoModelForImageTextToText, AutoProcessor
60
+ token = token or os.getenv("HF_TOKEN") or os.getenv("HUGGING_FACE_HUB_TOKEN")
61
+ dtype = torch.bfloat16 if str(device).startswith("cuda") else torch.float32
62
+ proc = AutoProcessor.from_pretrained(model_id, token=token)
63
+ model = AutoModelForImageTextToText.from_pretrained(
64
+ model_id, dtype=dtype, device_map=device, token=token)
65
+ model.eval()
66
+ for p in model.parameters():
67
+ p.requires_grad_(False)
68
+
69
+ # attach fast plastic LoRA on lm_head
70
+ head_name, head_mod = None, None
71
+ for name, mod in model.named_modules():
72
+ if name.endswith("lm_head") and isinstance(mod, nn.Linear):
73
+ head_name, head_mod = name, mod
74
+ head = _PlasticHead(head_mod, r=r_fast)
75
+ parent = model
76
+ *pth, last = head_name.split(".")
77
+ for pp in pth:
78
+ parent = getattr(parent, pp)
79
+ setattr(parent, last, head)
80
+ return cls(model, proc.tokenizer, head, **kw)
81
+
82
+ # ---------------- fast self-learning ----------------
83
+ def reset(self):
84
+ """Wipe fast adaptation -> exactly the strands-expert base again."""
85
+ import torch
86
+ import torch.nn as nn
87
+ with torch.no_grad():
88
+ self.head.B.zero_()
89
+ nn.init.normal_(self.head.A, std=0.01)
90
+ self.mean = None
91
+
92
+ def _nll(self, ids):
93
+ import torch
94
+ o = self.model(input_ids=ids)
95
+ lg = o.logits[:, :-1, :]
96
+ return torch.nn.functional.cross_entropy(
97
+ lg.reshape(-1, lg.size(-1)).float(), ids[:, 1:].reshape(-1))
98
+
99
+ def observe(self, text, learn=True, max_length=2048):
100
+ """Predict `text`; if surprised, rewrite the fast weights (bounded).
101
+ Returns the pre-update NLL (the surprise)."""
102
+ import torch
103
+ ids = self.tok(text, return_tensors="pt", truncation=True,
104
+ max_length=max_length).input_ids.to(self.device)
105
+ if ids.shape[1] < 2:
106
+ return None
107
+ with torch.no_grad():
108
+ e = self._nll(ids).item()
109
+ fire = (self.mean is None) or (e > self.mean + self.k_gate * abs(self.mean))
110
+ self.mean = e if self.mean is None else self.beta * self.mean + (1 - self.beta) * e
111
+ if learn and fire:
112
+ loss = self._nll(ids)
113
+ self.opt.zero_grad()
114
+ loss.backward()
115
+ torch.nn.utils.clip_grad_norm_([self.head.A, self.head.B], 1.0)
116
+ self.opt.step()
117
+ with torch.no_grad():
118
+ self.head.B.mul_(self.decay) # EMA decay = bounded plasticity
119
+ return e
120
+
121
+ # ---------------- chat ----------------
122
+ def chat(self, user_msg, max_new_tokens=512, temperature=0.7):
123
+ import torch
124
+ msgs = [{"role": "user", "content": user_msg}]
125
+ ids = self.tok.apply_chat_template(
126
+ msgs, add_generation_prompt=True, return_tensors="pt").to(self.device)
127
+ with torch.no_grad():
128
+ out = self.model.generate(
129
+ input_ids=ids, max_new_tokens=max_new_tokens,
130
+ do_sample=temperature > 0, temperature=max(temperature, 1e-5),
131
+ pad_token_id=self.tok.eos_token_id)
132
+ return self.tok.decode(out[0, ids.shape[1]:], skip_special_tokens=True)
133
+
134
+
135
+ def _plastic_head_cls():
136
+ import torch
137
+ import torch.nn as nn
138
+
139
+ class PlasticHead(nn.Module):
140
+ """Fast LoRA on lm_head: y = base(x) + scale*(x A)B. Only A,B change."""
141
+ def __init__(self, base, r=16, scale=2.0):
142
+ super().__init__()
143
+ self.base = base
144
+ for p in base.parameters():
145
+ p.requires_grad_(False)
146
+ dev, dt = base.weight.device, base.weight.dtype
147
+ self.A = nn.Parameter(torch.randn(base.in_features, r, device=dev, dtype=dt) * 0.01)
148
+ self.B = nn.Parameter(torch.zeros(r, base.out_features, device=dev, dtype=dt))
149
+ self.scale = scale
150
+
151
+ def forward(self, x):
152
+ return self.base(x) + self.scale * ((x @ self.A) @ self.B)
153
+
154
+ return PlasticHead
155
+
156
+
157
+ def _PlasticHead(base, r=16, scale=2.0):
158
+ return _plastic_head_cls()(base, r=r, scale=scale)