cccmmd commited on
Commit
9d0d4e9
·
1 Parent(s): 90cee57

feat: add Tiny-NLA activation explanation with trained model weights

Browse files

- Train Tiny-NLA (AV + AR) on Qwen3-0.6B to explain layer-19 activations
- Add /api/activation-explain endpoint with TinyNLAEngine (lazy loading)
- Integrate activation explain into Logit Lens frontend page
- Include training scripts, eval tools, and experiment configs
- Ship trained checkpoints: AV LoRA adapter + AR head (~20MB)

Files changed (45) hide show
  1. .gitignore +4 -0
  2. DESIGN_PROMPT.md +608 -0
  3. artifacts/tiny_nla/checkpoints/ar/best_ar_head.pt +3 -0
  4. artifacts/tiny_nla/checkpoints/av/README.md +207 -0
  5. artifacts/tiny_nla/checkpoints/av/adapter_config.json +43 -0
  6. artifacts/tiny_nla/checkpoints/av/adapter_model.safetensors +3 -0
  7. artifacts/tiny_nla/checkpoints/av/added_tokens.json +28 -0
  8. artifacts/tiny_nla/checkpoints/av/chat_template.jinja +89 -0
  9. artifacts/tiny_nla/checkpoints/av/merges.txt +0 -0
  10. artifacts/tiny_nla/checkpoints/av/special_tokens_map.json +25 -0
  11. artifacts/tiny_nla/checkpoints/av/tokenizer.json +3 -0
  12. artifacts/tiny_nla/checkpoints/av/tokenizer_config.json +239 -0
  13. artifacts/tiny_nla/checkpoints/av/vocab.json +0 -0
  14. backend/api/activation_explain.py +100 -0
  15. backend/core/tiny_nla.py +188 -0
  16. backend/platform/source_page.py +1 -1
  17. client/src/css/pages/logit_lens.scss +48 -0
  18. client/src/logit_lens.html +10 -0
  19. client/src/pages/logit_lens/index.ts +61 -1
  20. client/src/shared/api/GLTR_API.ts +55 -0
  21. client/src/shared/lang/translations.ts +5 -0
  22. experiments/tiny_nla/PLAN_v2.md +406 -0
  23. experiments/tiny_nla/batch_teacher_labels.py +105 -0
  24. experiments/tiny_nla/eval_final.py +179 -0
  25. experiments/tiny_nla/eval_roundtrip.py +157 -0
  26. experiments/tiny_nla/expand_dataset.py +82 -0
  27. experiments/tiny_nla/expand_dataset_v2.py +61 -0
  28. experiments/tiny_nla/gen_demo_html.py +43 -0
  29. experiments/tiny_nla/gen_opencode_only.py +99 -0
  30. experiments/tiny_nla/generate_data.py +397 -0
  31. experiments/tiny_nla/generate_teacher_labels.py +178 -0
  32. experiments/tiny_nla/generate_teacher_labels_opencode.py +207 -0
  33. experiments/tiny_nla/generate_teacher_labels_v2.py +186 -0
  34. experiments/tiny_nla/infer_tiny_nla.py +316 -0
  35. experiments/tiny_nla/nla_meta.yaml +48 -0
  36. experiments/tiny_nla/postprocess_dataset.py +112 -0
  37. experiments/tiny_nla/smoke_stage0.py +401 -0
  38. experiments/tiny_nla/train_ar.py +301 -0
  39. experiments/tiny_nla/train_av.py +331 -0
  40. experiments/tiny_nla/train_rl_grpo.py +404 -0
  41. experiments/tiny_nla/train_sft_v2.py +451 -0
  42. requirements.txt +1 -0
  43. server.py +1 -0
  44. server.yaml +61 -0
  45. train +394 -0
.gitignore CHANGED
@@ -8,6 +8,7 @@ __pycache__/
8
  *.py[cod]
9
  .venv/
10
  venv/
 
11
  env/
12
  node_modules/
13
  client/src/node_modules/
@@ -34,6 +35,9 @@ user_dialog_history/*
34
  .cache_huggingface/
35
  .env
36
 
 
 
 
37
  # --- 系统与 IDE ---
38
  .DS_Store
39
  .idea/
 
8
  *.py[cod]
9
  .venv/
10
  venv/
11
+ venv*/
12
  env/
13
  node_modules/
14
  client/src/node_modules/
 
35
  .cache_huggingface/
36
  .env
37
 
38
+ # --- 实验工件(模型权重、数据、检查点) ---
39
+ artifacts/
40
+
41
  # --- 系统与 IDE ---
42
  .DS_Store
43
  .idea/
DESIGN_PROMPT.md ADDED
@@ -0,0 +1,608 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # InfoLens / INNERSCOPE — 统一工作流重设计 Prompt
2
+
3
+ > **目标受众**:专精 UI/UX 设计的 Agent
4
+ > **任务**:理解现有前端架构和各功能模块,设计一个统一的工作流 demo,将当前孤立的 7 个功能(含新训练的 Tiny-NLA)有机融合为一个流畅的「输入文本 → 多维度探索模型内部」体验。
5
+ > **输出形式**:给出初步的 HTML/CSS/JS demo(可以是静态原型或可交互 demo),不需要对接真实后端 API。
6
+ > **重要的**:这是概念验证 demo,不是产品实现。你可以自由选择任何设计方向、技术栈、视觉风格。不要被现有代码的样式束缚。
7
+
8
+ ---
9
+
10
+ ## 一、项目是什么
11
+
12
+ **InfoLens(品牌名 INNERSCOPE)** 是一个「透视大语言模型内部」的可视化工具箱。用户输入一段文本,通过不同的分析维度观察模型在想什么。
13
+
14
+ - **服务端**:Python Flask/Connexion,跑在本地 M4 Pro Mac 上
15
+ - **模型**:Qwen3-0.6B (base + instruct),本地 CPU/MPS 推理
16
+ - **用户画像**:AI 研究者 / 可解释性爱好者,希望直观看到模型内部状态
17
+
18
+ ---
19
+
20
+ ## 二、现有前端代码结构(供参考,不需要遵循)
21
+
22
+ ```
23
+ client/src/
24
+ ├── index.html # 首页(导航网格)
25
+ ├── analysis.html # Info Highlight 页
26
+ ├── chat.html # Raw Chat 页
27
+ ├── attribution.html # Attribution 页
28
+ ├── causal_flow.html # Causal Flow 页
29
+ ├── logit_lens.html # Logit Lens 页
30
+ ├── branch_tree.html # Branch Tree 页
31
+
32
+ ├── webpack.config.js # 多入口打包
33
+
34
+ ├── pages/ # 每个页面的 TS 入口
35
+ │ ├── home/index.ts
36
+ │ ├── analysis/index.ts
37
+ │ ├── chat/index.ts # ~800 行,最复杂
38
+ │ ├── attribution/index.ts
39
+ │ ├── causal_flow/index.ts
40
+ │ ├── logit_lens/index.ts
41
+ │ └── branch_tree/index.ts
42
+
43
+ ├── shared/ # 跨页面共享基础设施
44
+ │ ├── api/GLTR_API.ts # TextAnalysisAPI 类(所有后端通信)
45
+ │ ├── core/ # URL 解析、事件总线、工具函数
46
+ │ ├── ui/ # toast, dialog, 主题/语言切换, 面板布局
47
+ │ ├── vis/ # D3 可视化组件(Token着色、Tooltip、直方图、散点图)
48
+ │ └── lang/translations.ts # i18n 翻译
49
+
50
+ ├── partials/ # HTML 片段(被 include 到页面中)
51
+
52
+ ├── features/ # 按功能域拆分
53
+ │ ├── analysis/ # Info Highlight 流程
54
+ │ ├── chat/ # 对话/续写逻辑
55
+ │ ├── causal_flow/ # Causal Flow 逻辑
56
+ │ └── demo/ # Demo 存储
57
+
58
+ └── css/pages/ # 每页独立的 SCSS
59
+ ```
60
+
61
+ ---
62
+
63
+ ## 三、7 个功能模块详解(含完整 API Schema)
64
+
65
+ ### 3.1 Info Highlight — Token 级别信息密度分析
66
+
67
+ **做什么**:展示模型对文本中每个 token 的「意外程度」——哪些词模型觉得理所当然,哪些词让它意外。
68
+
69
+ **API:`POST /api/analyze`**
70
+ ```
71
+ 请求:
72
+ {
73
+ "model": "base" | "instruct",
74
+ "text": "water freezes into solid ice",
75
+ "stream": false // 可选,true 时 SSE 流式返回
76
+ }
77
+
78
+ 响应 200:
79
+ {
80
+ "request": { "text": "water freezes into solid ice" },
81
+ "result": {
82
+ "model": "qwen3-0.6b-base",
83
+ "bpe_strings": [
84
+ {
85
+ "offset": [0, 5], // 字符偏移 [start, end)
86
+ "raw": "water", // token 原文
87
+ "real_topk": [156, 0.023], // [模型排序名次, softmax 概率]
88
+ "pred_topk": [ // 该位置 top-N 候选
89
+ ["ice", 0.34],
90
+ ["liquid", 0.18],
91
+ ...
92
+ ]
93
+ },
94
+ ...
95
+ ]
96
+ }
97
+ }
98
+
99
+ 错误:400 (缺参数) / 404 (模型未注册) / 500 (推理失败) / 503 (模型加载失败)
100
+ ```
101
+
102
+ **附加子功能 — 语义搜索:`POST /api/analyze-semantic`**
103
+ ```
104
+ 请求:
105
+ {
106
+ "query": "frozen", // 查询主题
107
+ "text": "...", // 原文
108
+ "stream": false, // 可选 SSE 流式
109
+ "submode": "count", // 可选: count / fill_blank
110
+ "debug_info": true // 可选,返回 top10 预测
111
+ }
112
+
113
+ 响应 200:
114
+ {
115
+ "success": true,
116
+ "model": "qwen3-0.6b-instruct",
117
+ "token_attention": [
118
+ {
119
+ "offset": [13, 19],
120
+ "raw": "solid",
121
+ "score": 0.87 // 对 query 的语义关注度
122
+ },
123
+ ...
124
+ ],
125
+ "debug_info": { // 仅 debug_info=true 时
126
+ "abbrev": "...",
127
+ "topk_tokens": [...],
128
+ "topk_probs": [...]
129
+ }
130
+ }
131
+ ```
132
+
133
+ ---
134
+
135
+ ### 3.2 Raw Chat — 对话 & 续写
136
+
137
+ **做什么**:和模型对话,可视化每次续写中每个生成 token 的细节。支持 Raw 模式(直接给 prompt)和 Chat 模式(system + user prompt + chat template)。
138
+
139
+ **API 链(3 步):**
140
+
141
+ **Step 1 — 组装 prompt:`POST /v1/completions/prompt`**
142
+ ```
143
+ 请求:
144
+ {
145
+ "model": "qwen3-0.6b-instruct",
146
+ "messages": [
147
+ { "role": "system", "content": "You are a helpful assistant." },
148
+ { "role": "user", "content": "What is water?" }
149
+ ],
150
+ "tools": [...], // 可选 OpenAI tools schema
151
+ "enable_thinking": false // 可选,启用 Qwen3 thinking 模式
152
+ }
153
+
154
+ 响应 200:
155
+ {
156
+ "prompt_used": "<|im_start|>system\n...<|im_end|>\n<|im_start|>user\nWhat is water?<|im_end|>\n<|im_start|>assistant\n"
157
+ }
158
+ ```
159
+
160
+ **Step 2 — 续写(SSE 流式):`POST /v1/completions`**
161
+ ```
162
+ 请求:
163
+ {
164
+ "model": "qwen3-0.6b-instruct",
165
+ "prompt": "<|im_start|>system\n...", // 从上一步或直接构造
166
+ "max_tokens": 256, // 可选,正整数
167
+ "temperature": 0.7, // 可选,0-2
168
+ "top_p": 0.9, // 可选,0-1
169
+ "stop": ["\n\n"] // 可选,停止序列(最多4个)
170
+ }
171
+ // 响应恒为 text/event-stream (SSE)
172
+ // type=delta → 增量 token
173
+ // type=result → 末条,data 同 OpenAICompletionsResponse
174
+
175
+ SSE 末条 data (type=result):
176
+ {
177
+ "id": "cmpl-xxx",
178
+ "object": "text_completion",
179
+ "created": 1719000000,
180
+ "model": "qwen3-0.6b-instruct",
181
+ "choices": [{
182
+ "text": "Water is a chemical substance...",
183
+ "index": 0,
184
+ "finish_reason": "stop" // stop / length / content_filter
185
+ }],
186
+ "usage": {
187
+ "prompt_tokens": 24,
188
+ "completion_tokens": 45,
189
+ "total_tokens": 69
190
+ },
191
+ "info_radar": { // 续写 token 级分析
192
+ "bpe_strings": [
193
+ { "offset": [0, 5], "raw": "Water", "real_topk": [1, 0.92], "pred_topk": [...] },
194
+ ...
195
+ ]
196
+ }
197
+ }
198
+ ```
199
+
200
+ **Step 3 — 停止生成:`POST /v1/completions/stop`**
201
+ ```
202
+ 请求:(无 body)
203
+ 响应 200:{ "ok": true }
204
+ ```
205
+
206
+ **高级功能**:Tool Calling 多轮模拟(通过 `/v1/completions/prompt-incremental` 计算 tool response 增量后缀)、Teacher Forcing(强制续写特定文本)、multi-turn 对话缓存与恢复。
207
+
208
+ ---
209
+
210
+ ### 3.3 Attribution — 输入对预测的因果归因
211
+
212
+ **做什么**:给定上下文和目标预测,分析「是哪些输入词驱动了这个预测」。提供两种方法:梯度归因(gradient-based)和消融归因(ablation/occlusion-based)。
213
+
214
+ **方法 A — 梯度归因:`POST /api/prediction-attribute`**
215
+ ```
216
+ 请求:
217
+ {
218
+ "context": "water freezes and turns into solid", // 必填,token 数 ≤ 2000
219
+ "target_prediction": " ice", // 可选,缺省用 top-1
220
+ "model": "base" | "instruct", // 必填
221
+ "source_page": "attribution", // 必填
222
+ "flow_id": null, // 可选,连续归因会话 ID
223
+ "flow_step": null // 可选,连续归因步骤
224
+ }
225
+
226
+ 响应 200:
227
+ {
228
+ "success": true,
229
+ "model": "qwen3-0.6b-base",
230
+ "target_token": "ice", // 归因目标 token
231
+ "target_prob": 0.34, // 预测概率
232
+ "token_attribution": [
233
+ {
234
+ "offset": [0, 5], // 字符偏移 [start, end]
235
+ "raw": "water",
236
+ "score": 0.42 // 梯度 L2 范数归因分(正值=支撑,越低越无关)
237
+ },
238
+ { "offset": [30, 35], "raw": "solid", "score": 0.78 },
239
+ ...
240
+ ],
241
+ "debug_info": { // 下一 token 的 top-10
242
+ "topk_tokens": ["ice", "liquid", "crystal", ...],
243
+ "topk_probs": [0.34, 0.18, 0.12, ...]
244
+ },
245
+ "is_eos": false // target 是否为 EOS token
246
+ }
247
+ ```
248
+
249
+ **方法 B — 消融归因:`POST /api/ablation-attribute`**
250
+ ```
251
+ 请求:
252
+ {
253
+ "context": "water freezes and turns into solid", // 必填,token 数 ≤ 500
254
+ "target_prediction": " ice", // 可选,与 target_token_id 互斥
255
+ "target_token_id": null, // 可选,直接指定 token id
256
+ "model": "base" | "instruct",
257
+ "source_page": "attribution"
258
+ }
259
+
260
+ 响应 200:
261
+ {
262
+ "success": true,
263
+ "target_token": "ice",
264
+ "target_prob": 0.34, // baseline 概率
265
+ "token_attribution": [
266
+ {
267
+ "offset": [0, 5],
268
+ "raw": "water",
269
+ "score": 0.12, // ΔP = baseline_prob − occluded_prob(可为负)
270
+ "delta_logit": 0.38 // Δlogit
271
+ },
272
+ ...
273
+ ],
274
+ "is_eos": false
275
+ }
276
+
277
+ 错误:400 (缺字段/model非法/超长) / 500 (推理失败) / 503 (繁忙)
278
+ ```
279
+
280
+ ---
281
+
282
+ ### 3.4 Logit Lens — 逐层解码轨迹
283
+
284
+ **做什么**:看模型在每一层的「中间想法」——把每层 hidden state 投影到词表(final norm + lm_head),展示每层 top-k 候选词和目标词概率如何逐层演化。
285
+
286
+ **API:`POST /api/logit-lens`**
287
+ ```
288
+ 请求:
289
+ {
290
+ "context": "The sun rises in the east and sets in the", // 必填,token 数 ≤ 500
291
+ "target_prediction": " west", // 可选,与 target_token_id 互斥
292
+ "target_token_id": null, // 可选
293
+ "model": "base" | "instruct",
294
+ "source_page": "logit_lens"
295
+ }
296
+
297
+ 响应 200:
298
+ {
299
+ "success": true,
300
+ "model": "qwen3-0.6b-base",
301
+ "target_token": "west",
302
+ "n_layers": 28, // Transformer 层数(不含 embedding)
303
+ "final_target_prob": 0.72, // 最终层目标 token 概率
304
+ "layers": [
305
+ {
306
+ "layer": 0,
307
+ "is_embedding": true, // 第 0 层为 embedding 层
308
+ "topk_tokens": ["the", "a", "in", "and", ...],
309
+ "topk_probs": [0.12, 0.09, 0.07, 0.06, ...],
310
+ "target_prob": 0.001 // 目标词在该层的概率
311
+ },
312
+ {
313
+ "layer": 1,
314
+ "is_embedding": false,
315
+ "topk_tokens": ["in", "the", "to", "of", ...],
316
+ "topk_probs": [0.15, 0.11, 0.08, 0.07, ...],
317
+ "target_prob": 0.003
318
+ },
319
+ ... // 共 n_layers+1 层(含 embedding)
320
+ {
321
+ "layer": 28,
322
+ "is_embedding": false,
323
+ "topk_tokens": ["west", "east", "north", ...], // 最终层
324
+ "topk_probs": [0.72, 0.08, 0.03, ...],
325
+ "target_prob": 0.72
326
+ }
327
+ ],
328
+ "is_eos": false
329
+ }
330
+ ```
331
+
332
+ **关键概念 — Eureka 层**:目标词首次跃至 top-1 的那一层。例如 layer 18 时 `west` 概率首次超过所有其他候选,这就是「模型想通了这是 west 的那一层」。
333
+
334
+ ---
335
+
336
+ ### 3.5 Branch Tree — 续写分支展开
337
+
338
+ **做什么**:输入前缀,展开模型所有可能的续写路径(top-k tree),递归探索「模型会怎么往下写」。
339
+
340
+ **API:`POST /api/branch-next`**
341
+ ```
342
+ 请求:
343
+ {
344
+ "prefix": "The sun rises in the", // 必填
345
+ "model": "base" | "instruct", // 必填
346
+ "source_page": "branch_tree", // 必填
347
+ "top_k": 10 // 可选,候选数(默认10,上限50)
348
+ }
349
+
350
+ 响应 200:
351
+ {
352
+ "success": true,
353
+ "model": "qwen3-0.6b-base",
354
+ "prefix_tokens": 5, // prefix 的 token 数
355
+ "candidates": [
356
+ { "token": "east", "token_id": 8423, "prob": 0.48 },
357
+ { "token": "west", "token_id": 2534, "prob": 0.22 },
358
+ { "token": "sky", "token_id": 6712, "prob": 0.08 },
359
+ ...
360
+ ],
361
+ "is_context_full": false // true 时前端应禁用展开
362
+ }
363
+
364
+ 错误:400 (缺参数/model非法/超长) / 500 / 503
365
+ ```
366
+
367
+ **递归模式**:选择一个候选 → 拼接进 prefix → 再次调用 branch-next → 形成树。前端负责树的构建和渲染逻辑。
368
+
369
+ ---
370
+
371
+ ### 3.6 Causal Flow — 因果信息流
372
+
373
+ **做什么**:追踪信息如何在模型层间流动,观察归因信号沿生成步骤的传播。核心是「文本自回归生成 + 每一步的输入 token 对输出 token 的梯度归因」的串行管道。
374
+
375
+ **依赖的 API**(Causal Flow 是客户端编排的复合流程):
376
+ 1. `POST /api/tokenize` — 分词(见下方 3.6.1)
377
+ 2. `POST /api/branch-next` — 获取 top-k 候选(同 3.5)
378
+ 3. `POST /api/prediction-attribute` — 每一步的梯度归因(同 3.3),带 `flow_id` + `flow_step` 连续归因
379
+ 4. `POST /v1/completions/prompt` — 组装 prompt(同 3.2 Step 1)
380
+ 5. `POST /v1/completions` — 续写生成下一个 token
381
+
382
+ **客户端流程**:输入 context → tokenize → 自动循环:用当前 context 调用 branch-next,选 top-1 → 调用 prediction-attribute 归因当前步 → 拼接 token → 继续下一步(或手动选择分支方向)
383
+
384
+ **预计算 Demo 数据**:可加载预计算的完整归因 DAG,以 D3 力导向图或传播动画播放,无需实时推理。
385
+
386
+ ---
387
+
388
+ ### 3.6.1 Tokenize — 文本分词(通用工具端点)
389
+
390
+ **API:`POST /api/tokenize`**
391
+ ```
392
+ 请求:
393
+ {
394
+ "context": "Hello, world!", // 必填
395
+ "model": "base" | "instruct" // 必填
396
+ }
397
+
398
+ 响应 200:
399
+ {
400
+ "success": true,
401
+ "spans": [
402
+ { "offset": [0, 5], "raw": "Hello" },
403
+ { "offset": [5, 6], "raw": "," },
404
+ { "offset": [6, 7], "raw": " " },
405
+ { "offset": [7, 12], "raw": "world" },
406
+ { "offset": [12, 13], "raw": "!" }
407
+ ]
408
+ }
409
+
410
+ 错误:400 (缺字段/model非法)
411
+ ```
412
+
413
+ **特点**:不持有推理锁,不做前向/梯度计算,响应极快。是构建统一工作流的「第一步」。
414
+
415
+ ---
416
+
417
+ ### 3.7 🆕 Tiny-NLA Activation Explainer(新训练的模型)
418
+
419
+ **做什么**:给定一个 token 位置的激活向量(从模型 layer 19 残差流提取的 1024 维向量),用 LoRA 微调的小模型解释「这个激活向量代表了什么概念」。支持两种输入模式:传入文本让后端自动提取激活,或直接传入向量。
420
+
421
+ **模型细节**:
422
+ - **架构**:Qwen3-0.6B-Base + LoRA adapter (r=8, alpha=16) @ layer 19 + AR head(线性层 1024→1024)
423
+ - **推理精度**:float32 + eager(独立加载,不复用 base 槽位的 float16)
424
+ - **训练数据**:284 条 teacher labeled 激活-解释对
425
+ - **设备**:Apple M4 Pro MPS
426
+ - **核心机制**:用特殊 token `㈎`(id=149705)作为激活注入锚点,注入归一化后的激活向量(injection_scale=126.223),AV 模型 generate 出自然语言解释,AR head 重建后计算 cosine 评估可信度
427
+
428
+ **API:`POST /api/activation-explain`**
429
+ ```
430
+ 请求模式 A — 文本模式(自动提取激活):
431
+ {
432
+ "model": "base" | "instruct", // 必填
433
+ "source_page": "logit_lens", // 必填
434
+ "text": "water freezes into solid ice", // 输入上下文
435
+ "token_index": 4 // 目标 token 位置(text 模式必填,必须 ≥ 0 且在范围内)
436
+ }
437
+
438
+ 请求模式 B — 向量模式(直接传入):
439
+ {
440
+ "model": "base",
441
+ "source_page": "workspace",
442
+ "vector": [0.12, -0.34, 0.56, ...] // 恰好 1024 个 float(与 text 互斥)
443
+ }
444
+
445
+ 响应 200:
446
+ {
447
+ "success": true,
448
+ "concept": "", // 概念标签(当前为空,可后续扩展)
449
+ "explanation": "This activation vector encodes the concept of water turning into a solid state through freezing, capturing the physical phase transition from liquid to ice.",
450
+ "roundtrip_cosine": 0.6631, // 重建可信度分数 (0-1)
451
+ "vector_dim": 1024, // 向量维度
452
+ "note": "" // 备注或状态说明
453
+ }
454
+
455
+ 错误:
456
+ 400 — 缺字段、model 非法、text/vector 均未提供、token_index 越界、向量维度不为 1024
457
+ 500 — Tiny-NLA 推理失败
458
+ 503 — 推理锁等待超时(30s)
459
+ ```
460
+
461
+ **可信度解读**(roundtrip_cosine):
462
+ - **≥ 0.70** — 优秀(Excellent):解释高度忠实于激活
463
+ - **0.60–0.70** — 良好(Good):解释与激活有较强关联
464
+ - **0.50–0.60** — 可用(Fair):解释方向大致正确
465
+ - **< 0.50** — 保留(Low):解释与激活关联较弱,仅供参考
466
+
467
+ **模型性能基线**:
468
+ - AV(Activation Vector 模型):best_val_loss = 0.64
469
+ - AR(Auto-Regressive head):best_val_cosine = 0.6631,mean baseline cosine = 0.5962,shuffled baseline = 0.5252
470
+ - 相比随机 baseline 提升 0.07(统计显著)
471
+
472
+ ## 四、核心问题
473
+
474
+ ```
475
+ 现状:6 个功能 = 6 个独立页面
476
+ ❌ 同一段文本要在不同页面反复输入
477
+ ❌ 结果之间没有关联(Logit Lens 的结果不能顺手看 Attribution)
478
+ ❌ 新模型(Tiny-NLA)完全没接入前端
479
+ ❌ 用户体验割裂——每个功能像独立的工具而非一个整体
480
+ ```
481
+
482
+ ---
483
+
484
+ ## 五、设计目标
485
+
486
+ > **输入一次文本,在一个页面上以多种「透镜」/「视角」探索模型内部。**
487
+
488
+ 期望的用户旅程:
489
+ 1. 用户输入一段文本 → tokenize 为可交互的 token 序列
490
+ 2. 用户选中某个 token → 可以从多个维度探索:
491
+ - 📊 **Info Highlight**:看 surprisal
492
+ - 🎯 **Attribution**:看哪些前文驱动了这个 token
493
+ - 🔍 **Logit Lens**:看逐层解码轨迹
494
+ - 🌳 **Branch Tree**:展开可能的续写
495
+ - 🧠 **Activation Explainer**:解释激活向量代表什么概念
496
+ - 🔗 **Causal Flow**:看信息流向
497
+ 3. 多个分析视图可以同时打开,共享同一份输入上下文
498
+
499
+ **你可以完全自由地设计交互模式和视觉风格**,不需要考虑现有前端代码的任何约束。
500
+
501
+ ---
502
+
503
+ ## 六、Mock 数据 & 接口设计自由度
504
+
505
+ ### 现有接口参考
506
+
507
+ 所有后端的完整请求/响应 schema 见**第三章各小节**。demo 中构造 mock 数据时,请直接参考对应功能的 API 响应结构。
508
+
509
+ ### 🆕 可以设计新接口
510
+
511
+ **你不必拘泥于现有接口。** 如果统一工作流的交互需要新的后端接口来达到更好的效果,完全可以提出新的 API 设计。
512
+
513
+ **前提条件**:新接口必须在我们后端能力范围内。我们后端的能力边界是:
514
+
515
+ - **运行时**:Python Flask/Connexion,单个 `server.py` 注册函数即自动生成路由
516
+ - **模型能力**:
517
+ - Qwen3-0.6B (base + instruct) — 前向推理、logits 获取、hidden states 提取、tokenize
518
+ - Tiny-NLA (LoRA + AR head) — 激活解释(已有)
519
+ - **可做的计算**:
520
+ - 单个 token 的前向传播及 hidden states
521
+ - 任意层的 hidden state 投影到词表(logit lens 的核心操作)
522
+ - 梯度计算(需 backward)
523
+ - 两个向量的 cosine 相似度、L2 距离
524
+ - token 级别的概率获取
525
+ - **做不到的**:
526
+ - 大规模 batch 训练(但我们的训练代码在 `experiments/tiny_nla/` 中,可以离线跑)
527
+ - 需要额外预训练模型的计算(除非把模型文件放到 `artifacts/` 下)
528
+
529
+ **如果你设计了新接口,在 demo 中用 mock 数据即可。我们在后续实现阶段会在后端补上。**
530
+
531
+ 关键端点速查:
532
+ - **Info Highlight** → 3.1 节 `POST /api/analyze` 响应
533
+ - **Info Highlight 语义搜索** → 3.1 节 `POST /api/analyze-semantic` 响应
534
+ - **Raw Chat** → 3.2 节 `POST /v1/completions` SSE 末条响应
535
+ - **Attribution 梯度** → 3.3 节 `POST /api/prediction-attribute` 响应
536
+ - **Attribution 消融** → 3.3 节 `POST /api/ablation-attribute` 响应
537
+ - **Logit Lens** → 3.4 节 `POST /api/logit-lens` 响应(注意 layers 数组结构)
538
+ - **Branch Tree** → 3.5 节 `POST /api/branch-next` 响应
539
+ - **Causal Flow** → 3.6 节复合流程(tokenize + branch-next + prediction-attribute 串行)
540
+ - **Tokenize** → 3.6.1 节 `POST /api/tokenize` 响应
541
+ - **🆕 Activation Explainer** → 3.7 节 `POST /api/activation-explain` 响应
542
+
543
+ ---
544
+
545
+ ## 七、设计资源 & 参考
546
+
547
+ ### 设计灵感网站(自由参考)
548
+ - **https://21st.dev/** — 现代 UI 组件灵感和设计模式
549
+ - **https://cuicui.day/application-ui** — 应用 UI 设计参考
550
+ - **https://ui.aceternity.com/** — 创新 UI 组件和动画效果
551
+
552
+ ### 本地设计仓库(内含设计规范)
553
+
554
+ **[impeccable](https://github.com/pbakaus/impeccable)** — Anti-Design-Slop 检测器
555
+ - `AGENTS.md` — 仓库规范和工作流
556
+ - `DESIGN.md` — Neo Kinpaku 设计系统(暗色漆器 + 金箔 + 铜绿),包含完整的色彩、排版、圆角、间距 token 和组件规范(button、input、card、nav-link 等)
557
+ - `CLAUDE.md` — Claude 专用的详细行为规范和交互模式
558
+ - `docs/STYLE.md` — 风格指南
559
+ - **核心哲学**:避免 AI 生成设计中的常见反模式(AI-purple gradients、三列等大 Feature Cards、过度 glassmorphism、无限循环微动画等)
560
+
561
+ **[taste-skill](https://github.com/Leonxlnx/taste-skill)** — Anti-Slop 前端 Skill
562
+ - `skills/taste-skill/SKILL.md` — 核心设计 Skill(85KB),包含:
563
+ - 「Read the Room」——先推断设计意图再动手
564
+ - 「三旋钮」系统(DESIGN_VARIANCE / MOTION_INTENSITY / VISUAL_DENSITY)从 1-10 控制视觉风格
565
+ - 设计系统映射表(何时用 Fluent UI / Material / Carbon / Primer / govuk-frontend / shadcn 等)
566
+ - 反默认纪律(避免 LLM 默认审美)
567
+ - `skills/minimalist-skill/SKILL.md` — 极简风格
568
+ - `skills/brutalist-skill/SKILL.md` — 粗野主义风格
569
+ - `skills/soft-skill/SKILL.md` — 柔和风格
570
+ - `skills/redesign-skill/SKILL.md` — 重设计方法论
571
+ - `skills/stitch-skill/SKILL.md` — 设计缝合参考
572
+
573
+ **建议在开始设计前,至少阅读 `impeccable/DESIGN.md` 和 `taste-skill/skills/taste-skill/SKILL.md` 以了解高质量设计的核心原则。**
574
+
575
+ ---
576
+
577
+ ## 八、交付要求
578
+
579
+ ### 必须包含
580
+ 1. **统一输入区**:一个文本框接受用户输入,tokenize 后展示为可交互的 token 序列
581
+ 2. **至少 3 种分析入口**:包括至少 Logit Lens、Attribution、Tiny-NLA Activation Explainer
582
+ 3. **多视图共存**:可以同时看到至少 2 种分析的输出(不互斥)
583
+ 4. **Token 选择交互**:点击 token 触发对应分析
584
+ 5. **Mock 数据**:用静态数据代替真实 API 响应
585
+
586
+ ### 不需要做的
587
+ - 不需要对接真实后端 API
588
+ - 不需要实现完整的 i18n
589
+ - 不需要考虑 webpack 集成(独立 HTML 即可)
590
+ - 不需要实现 Admin / Settings 功能
591
+
592
+ ### 交付物
593
+ 一个或多个独立的 HTML/CSS/JS 文件,可以在浏览器中直接打开运行。
594
+
595
+ ---
596
+
597
+ ## 九、自由发挥空间
598
+
599
+ 以下内容**完全由你决定**,不受任何限制:
600
+
601
+ - 🎨 **视觉风格**:暗色/亮色、极简/丰富、任何配色方案
602
+ - 🖼️ **布局方式**:面板/分屏/Tab/卡片网格/画布/任何创意布局
603
+ - 🎬 **动效设计**:过渡/微交互/动画,完全自由
604
+ - 📐 **交互模式**:点击/悬停/拖拽/快捷键/手势
605
+ - 🏗️ **技术选型**:原生 CSS / Tailwind / 任何 CSS 框架 / 任何 JS 库
606
+ - 📝 **文案风格**:技术性 / 产品化 / 趣味化,任意
607
+
608
+ **唯一的要求是:看起来不像"又一个 AI 生成的 demo",而是像一个真正有设计思考的产品。**
artifacts/tiny_nla/checkpoints/ar/best_ar_head.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:af5e2021a16d78fce920c99559f4202e929ab7116b01d0c37c4986ffca3dbc1c
3
+ size 4195637
artifacts/tiny_nla/checkpoints/av/README.md ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ base_model: Qwen/Qwen3-0.6B
3
+ library_name: peft
4
+ pipeline_tag: text-generation
5
+ tags:
6
+ - base_model:adapter:Qwen/Qwen3-0.6B
7
+ - lora
8
+ - transformers
9
+ ---
10
+
11
+ # Model Card for Model ID
12
+
13
+ <!-- Provide a quick summary of what the model is/does. -->
14
+
15
+
16
+
17
+ ## Model Details
18
+
19
+ ### Model Description
20
+
21
+ <!-- Provide a longer summary of what this model is. -->
22
+
23
+
24
+
25
+ - **Developed by:** [More Information Needed]
26
+ - **Funded by [optional]:** [More Information Needed]
27
+ - **Shared by [optional]:** [More Information Needed]
28
+ - **Model type:** [More Information Needed]
29
+ - **Language(s) (NLP):** [More Information Needed]
30
+ - **License:** [More Information Needed]
31
+ - **Finetuned from model [optional]:** [More Information Needed]
32
+
33
+ ### Model Sources [optional]
34
+
35
+ <!-- Provide the basic links for the model. -->
36
+
37
+ - **Repository:** [More Information Needed]
38
+ - **Paper [optional]:** [More Information Needed]
39
+ - **Demo [optional]:** [More Information Needed]
40
+
41
+ ## Uses
42
+
43
+ <!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
44
+
45
+ ### Direct Use
46
+
47
+ <!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
48
+
49
+ [More Information Needed]
50
+
51
+ ### Downstream Use [optional]
52
+
53
+ <!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
54
+
55
+ [More Information Needed]
56
+
57
+ ### Out-of-Scope Use
58
+
59
+ <!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
60
+
61
+ [More Information Needed]
62
+
63
+ ## Bias, Risks, and Limitations
64
+
65
+ <!-- This section is meant to convey both technical and sociotechnical limitations. -->
66
+
67
+ [More Information Needed]
68
+
69
+ ### Recommendations
70
+
71
+ <!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
72
+
73
+ Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
74
+
75
+ ## How to Get Started with the Model
76
+
77
+ Use the code below to get started with the model.
78
+
79
+ [More Information Needed]
80
+
81
+ ## Training Details
82
+
83
+ ### Training Data
84
+
85
+ <!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
86
+
87
+ [More Information Needed]
88
+
89
+ ### Training Procedure
90
+
91
+ <!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
92
+
93
+ #### Preprocessing [optional]
94
+
95
+ [More Information Needed]
96
+
97
+
98
+ #### Training Hyperparameters
99
+
100
+ - **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
101
+
102
+ #### Speeds, Sizes, Times [optional]
103
+
104
+ <!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
105
+
106
+ [More Information Needed]
107
+
108
+ ## Evaluation
109
+
110
+ <!-- This section describes the evaluation protocols and provides the results. -->
111
+
112
+ ### Testing Data, Factors & Metrics
113
+
114
+ #### Testing Data
115
+
116
+ <!-- This should link to a Dataset Card if possible. -->
117
+
118
+ [More Information Needed]
119
+
120
+ #### Factors
121
+
122
+ <!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
123
+
124
+ [More Information Needed]
125
+
126
+ #### Metrics
127
+
128
+ <!-- These are the evaluation metrics being used, ideally with a description of why. -->
129
+
130
+ [More Information Needed]
131
+
132
+ ### Results
133
+
134
+ [More Information Needed]
135
+
136
+ #### Summary
137
+
138
+
139
+
140
+ ## Model Examination [optional]
141
+
142
+ <!-- Relevant interpretability work for the model goes here -->
143
+
144
+ [More Information Needed]
145
+
146
+ ## Environmental Impact
147
+
148
+ <!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
149
+
150
+ Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
151
+
152
+ - **Hardware Type:** [More Information Needed]
153
+ - **Hours used:** [More Information Needed]
154
+ - **Cloud Provider:** [More Information Needed]
155
+ - **Compute Region:** [More Information Needed]
156
+ - **Carbon Emitted:** [More Information Needed]
157
+
158
+ ## Technical Specifications [optional]
159
+
160
+ ### Model Architecture and Objective
161
+
162
+ [More Information Needed]
163
+
164
+ ### Compute Infrastructure
165
+
166
+ [More Information Needed]
167
+
168
+ #### Hardware
169
+
170
+ [More Information Needed]
171
+
172
+ #### Software
173
+
174
+ [More Information Needed]
175
+
176
+ ## Citation [optional]
177
+
178
+ <!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
179
+
180
+ **BibTeX:**
181
+
182
+ [More Information Needed]
183
+
184
+ **APA:**
185
+
186
+ [More Information Needed]
187
+
188
+ ## Glossary [optional]
189
+
190
+ <!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
191
+
192
+ [More Information Needed]
193
+
194
+ ## More Information [optional]
195
+
196
+ [More Information Needed]
197
+
198
+ ## Model Card Authors [optional]
199
+
200
+ [More Information Needed]
201
+
202
+ ## Model Card Contact
203
+
204
+ [More Information Needed]
205
+ ### Framework versions
206
+
207
+ - PEFT 0.19.1
artifacts/tiny_nla/checkpoints/av/adapter_config.json ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "alora_invocation_tokens": null,
3
+ "alpha_pattern": {},
4
+ "arrow_config": null,
5
+ "auto_mapping": null,
6
+ "base_model_name_or_path": "Qwen/Qwen3-0.6B",
7
+ "bias": "none",
8
+ "corda_config": null,
9
+ "ensure_weight_tying": false,
10
+ "eva_config": null,
11
+ "exclude_modules": null,
12
+ "fan_in_fan_out": false,
13
+ "inference_mode": true,
14
+ "init_lora_weights": true,
15
+ "layer_replication": null,
16
+ "layers_pattern": null,
17
+ "layers_to_transform": null,
18
+ "loftq_config": {},
19
+ "lora_alpha": 16,
20
+ "lora_bias": false,
21
+ "lora_dropout": 0.1,
22
+ "lora_ga_config": null,
23
+ "megatron_config": null,
24
+ "megatron_core": "megatron.core",
25
+ "modules_to_save": null,
26
+ "peft_type": "LORA",
27
+ "peft_version": "0.19.1",
28
+ "qalora_group_size": 16,
29
+ "r": 8,
30
+ "rank_pattern": {},
31
+ "revision": null,
32
+ "target_modules": [
33
+ "q_proj",
34
+ "v_proj"
35
+ ],
36
+ "target_parameters": null,
37
+ "task_type": "CAUSAL_LM",
38
+ "trainable_token_indices": null,
39
+ "use_bdlora": null,
40
+ "use_dora": false,
41
+ "use_qalora": false,
42
+ "use_rslora": false
43
+ }
artifacts/tiny_nla/checkpoints/av/adapter_model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c2bb9c1a38d1b11a5e99ed9c770bd08a16b5baa4277478f2925b7d839dafef7b
3
+ size 4602248
artifacts/tiny_nla/checkpoints/av/added_tokens.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "</think>": 151668,
3
+ "</tool_call>": 151658,
4
+ "</tool_response>": 151666,
5
+ "<think>": 151667,
6
+ "<tool_call>": 151657,
7
+ "<tool_response>": 151665,
8
+ "<|box_end|>": 151649,
9
+ "<|box_start|>": 151648,
10
+ "<|endoftext|>": 151643,
11
+ "<|file_sep|>": 151664,
12
+ "<|fim_middle|>": 151660,
13
+ "<|fim_pad|>": 151662,
14
+ "<|fim_prefix|>": 151659,
15
+ "<|fim_suffix|>": 151661,
16
+ "<|im_end|>": 151645,
17
+ "<|im_start|>": 151644,
18
+ "<|image_pad|>": 151655,
19
+ "<|object_ref_end|>": 151647,
20
+ "<|object_ref_start|>": 151646,
21
+ "<|quad_end|>": 151651,
22
+ "<|quad_start|>": 151650,
23
+ "<|repo_name|>": 151663,
24
+ "<|video_pad|>": 151656,
25
+ "<|vision_end|>": 151653,
26
+ "<|vision_pad|>": 151654,
27
+ "<|vision_start|>": 151652
28
+ }
artifacts/tiny_nla/checkpoints/av/chat_template.jinja ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {%- if tools %}
2
+ {{- '<|im_start|>system\n' }}
3
+ {%- if messages[0].role == 'system' %}
4
+ {{- messages[0].content + '\n\n' }}
5
+ {%- endif %}
6
+ {{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
7
+ {%- for tool in tools %}
8
+ {{- "\n" }}
9
+ {{- tool | tojson }}
10
+ {%- endfor %}
11
+ {{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
12
+ {%- else %}
13
+ {%- if messages[0].role == 'system' %}
14
+ {{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
15
+ {%- endif %}
16
+ {%- endif %}
17
+ {%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
18
+ {%- for message in messages[::-1] %}
19
+ {%- set index = (messages|length - 1) - loop.index0 %}
20
+ {%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
21
+ {%- set ns.multi_step_tool = false %}
22
+ {%- set ns.last_query_index = index %}
23
+ {%- endif %}
24
+ {%- endfor %}
25
+ {%- for message in messages %}
26
+ {%- if message.content is string %}
27
+ {%- set content = message.content %}
28
+ {%- else %}
29
+ {%- set content = '' %}
30
+ {%- endif %}
31
+ {%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
32
+ {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
33
+ {%- elif message.role == "assistant" %}
34
+ {%- set reasoning_content = '' %}
35
+ {%- if message.reasoning_content is string %}
36
+ {%- set reasoning_content = message.reasoning_content %}
37
+ {%- else %}
38
+ {%- if '</think>' in content %}
39
+ {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
40
+ {%- set content = content.split('</think>')[-1].lstrip('\n') %}
41
+ {%- endif %}
42
+ {%- endif %}
43
+ {%- if loop.index0 > ns.last_query_index %}
44
+ {%- if loop.last or (not loop.last and reasoning_content) %}
45
+ {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
46
+ {%- else %}
47
+ {{- '<|im_start|>' + message.role + '\n' + content }}
48
+ {%- endif %}
49
+ {%- else %}
50
+ {{- '<|im_start|>' + message.role + '\n' + content }}
51
+ {%- endif %}
52
+ {%- if message.tool_calls %}
53
+ {%- for tool_call in message.tool_calls %}
54
+ {%- if (loop.first and content) or (not loop.first) %}
55
+ {{- '\n' }}
56
+ {%- endif %}
57
+ {%- if tool_call.function %}
58
+ {%- set tool_call = tool_call.function %}
59
+ {%- endif %}
60
+ {{- '<tool_call>\n{"name": "' }}
61
+ {{- tool_call.name }}
62
+ {{- '", "arguments": ' }}
63
+ {%- if tool_call.arguments is string %}
64
+ {{- tool_call.arguments }}
65
+ {%- else %}
66
+ {{- tool_call.arguments | tojson }}
67
+ {%- endif %}
68
+ {{- '}\n</tool_call>' }}
69
+ {%- endfor %}
70
+ {%- endif %}
71
+ {{- '<|im_end|>\n' }}
72
+ {%- elif message.role == "tool" %}
73
+ {%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
74
+ {{- '<|im_start|>user' }}
75
+ {%- endif %}
76
+ {{- '\n<tool_response>\n' }}
77
+ {{- content }}
78
+ {{- '\n</tool_response>' }}
79
+ {%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
80
+ {{- '<|im_end|>\n' }}
81
+ {%- endif %}
82
+ {%- endif %}
83
+ {%- endfor %}
84
+ {%- if add_generation_prompt %}
85
+ {{- '<|im_start|>assistant\n' }}
86
+ {%- if enable_thinking is defined and enable_thinking is false %}
87
+ {{- '<think>\n\n</think>\n\n' }}
88
+ {%- endif %}
89
+ {%- endif %}
artifacts/tiny_nla/checkpoints/av/merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
artifacts/tiny_nla/checkpoints/av/special_tokens_map.json ADDED
@@ -0,0 +1,25 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|object_ref_start|>",
6
+ "<|object_ref_end|>",
7
+ "<|box_start|>",
8
+ "<|box_end|>",
9
+ "<|quad_start|>",
10
+ "<|quad_end|>",
11
+ "<|vision_start|>",
12
+ "<|vision_end|>",
13
+ "<|vision_pad|>",
14
+ "<|image_pad|>",
15
+ "<|video_pad|>"
16
+ ],
17
+ "eos_token": {
18
+ "content": "<|im_end|>",
19
+ "lstrip": false,
20
+ "normalized": false,
21
+ "rstrip": false,
22
+ "single_word": false
23
+ },
24
+ "pad_token": "<|im_end|>"
25
+ }
artifacts/tiny_nla/checkpoints/av/tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
3
+ size 11422654
artifacts/tiny_nla/checkpoints/av/tokenizer_config.json ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": false,
3
+ "add_prefix_space": false,
4
+ "added_tokens_decoder": {
5
+ "151643": {
6
+ "content": "<|endoftext|>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "151644": {
14
+ "content": "<|im_start|>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "151645": {
22
+ "content": "<|im_end|>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "151646": {
30
+ "content": "<|object_ref_start|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "151647": {
38
+ "content": "<|object_ref_end|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "151648": {
46
+ "content": "<|box_start|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "151649": {
54
+ "content": "<|box_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "151650": {
62
+ "content": "<|quad_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "151651": {
70
+ "content": "<|quad_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "151652": {
78
+ "content": "<|vision_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ },
85
+ "151653": {
86
+ "content": "<|vision_end|>",
87
+ "lstrip": false,
88
+ "normalized": false,
89
+ "rstrip": false,
90
+ "single_word": false,
91
+ "special": true
92
+ },
93
+ "151654": {
94
+ "content": "<|vision_pad|>",
95
+ "lstrip": false,
96
+ "normalized": false,
97
+ "rstrip": false,
98
+ "single_word": false,
99
+ "special": true
100
+ },
101
+ "151655": {
102
+ "content": "<|image_pad|>",
103
+ "lstrip": false,
104
+ "normalized": false,
105
+ "rstrip": false,
106
+ "single_word": false,
107
+ "special": true
108
+ },
109
+ "151656": {
110
+ "content": "<|video_pad|>",
111
+ "lstrip": false,
112
+ "normalized": false,
113
+ "rstrip": false,
114
+ "single_word": false,
115
+ "special": true
116
+ },
117
+ "151657": {
118
+ "content": "<tool_call>",
119
+ "lstrip": false,
120
+ "normalized": false,
121
+ "rstrip": false,
122
+ "single_word": false,
123
+ "special": false
124
+ },
125
+ "151658": {
126
+ "content": "</tool_call>",
127
+ "lstrip": false,
128
+ "normalized": false,
129
+ "rstrip": false,
130
+ "single_word": false,
131
+ "special": false
132
+ },
133
+ "151659": {
134
+ "content": "<|fim_prefix|>",
135
+ "lstrip": false,
136
+ "normalized": false,
137
+ "rstrip": false,
138
+ "single_word": false,
139
+ "special": false
140
+ },
141
+ "151660": {
142
+ "content": "<|fim_middle|>",
143
+ "lstrip": false,
144
+ "normalized": false,
145
+ "rstrip": false,
146
+ "single_word": false,
147
+ "special": false
148
+ },
149
+ "151661": {
150
+ "content": "<|fim_suffix|>",
151
+ "lstrip": false,
152
+ "normalized": false,
153
+ "rstrip": false,
154
+ "single_word": false,
155
+ "special": false
156
+ },
157
+ "151662": {
158
+ "content": "<|fim_pad|>",
159
+ "lstrip": false,
160
+ "normalized": false,
161
+ "rstrip": false,
162
+ "single_word": false,
163
+ "special": false
164
+ },
165
+ "151663": {
166
+ "content": "<|repo_name|>",
167
+ "lstrip": false,
168
+ "normalized": false,
169
+ "rstrip": false,
170
+ "single_word": false,
171
+ "special": false
172
+ },
173
+ "151664": {
174
+ "content": "<|file_sep|>",
175
+ "lstrip": false,
176
+ "normalized": false,
177
+ "rstrip": false,
178
+ "single_word": false,
179
+ "special": false
180
+ },
181
+ "151665": {
182
+ "content": "<tool_response>",
183
+ "lstrip": false,
184
+ "normalized": false,
185
+ "rstrip": false,
186
+ "single_word": false,
187
+ "special": false
188
+ },
189
+ "151666": {
190
+ "content": "</tool_response>",
191
+ "lstrip": false,
192
+ "normalized": false,
193
+ "rstrip": false,
194
+ "single_word": false,
195
+ "special": false
196
+ },
197
+ "151667": {
198
+ "content": "<think>",
199
+ "lstrip": false,
200
+ "normalized": false,
201
+ "rstrip": false,
202
+ "single_word": false,
203
+ "special": false
204
+ },
205
+ "151668": {
206
+ "content": "</think>",
207
+ "lstrip": false,
208
+ "normalized": false,
209
+ "rstrip": false,
210
+ "single_word": false,
211
+ "special": false
212
+ }
213
+ },
214
+ "additional_special_tokens": [
215
+ "<|im_start|>",
216
+ "<|im_end|>",
217
+ "<|object_ref_start|>",
218
+ "<|object_ref_end|>",
219
+ "<|box_start|>",
220
+ "<|box_end|>",
221
+ "<|quad_start|>",
222
+ "<|quad_end|>",
223
+ "<|vision_start|>",
224
+ "<|vision_end|>",
225
+ "<|vision_pad|>",
226
+ "<|image_pad|>",
227
+ "<|video_pad|>"
228
+ ],
229
+ "bos_token": null,
230
+ "clean_up_tokenization_spaces": false,
231
+ "eos_token": "<|im_end|>",
232
+ "errors": "replace",
233
+ "extra_special_tokens": {},
234
+ "model_max_length": 131072,
235
+ "pad_token": "<|im_end|>",
236
+ "split_special_tokens": false,
237
+ "tokenizer_class": "Qwen2Tokenizer",
238
+ "unk_token": null
239
+ }
artifacts/tiny_nla/checkpoints/av/vocab.json ADDED
The diff for this file is too large to render. See raw diff
 
backend/api/activation_explain.py ADDED
@@ -0,0 +1,100 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Activation Explainer API (Tiny-NLA)"""
2
+ import gc
3
+ import time
4
+
5
+ import torch
6
+
7
+ from backend.platform.oom import exit_if_oom
8
+ from backend.api.analyze import LOCK_WAIT_TIMEOUT
9
+ from backend.platform.access_log import get_client_ip, log_prediction_attribute_request
10
+ from backend.platform.source_page import ALLOWED_SOURCE_PAGES, normalize_source_page
11
+ from backend.core.tiny_nla import TinyNLAEngine, tiny_nla_lock, TINY_NLA_LOCK_TIMEOUT
12
+
13
+
14
+ def activation_explain(activation_explain_request):
15
+ model = activation_explain_request.get("model")
16
+ source_page = activation_explain_request.get("source_page")
17
+ text = activation_explain_request.get("text")
18
+ token_index = activation_explain_request.get("token_index")
19
+ vector = activation_explain_request.get("vector")
20
+
21
+ if model is None:
22
+ return {"success": False, "message": "Missing required field: model"}, 400
23
+ if model not in ("base", "instruct"):
24
+ return {"success": False, "message": 'model must be "base" or "instruct"'}, 400
25
+
26
+ if source_page is None or source_page == "":
27
+ return {"success": False, "message": "Missing required field: source_page"}, 400
28
+ normalized_source_page = normalize_source_page(source_page)
29
+ if normalized_source_page is None:
30
+ allowed = ", ".join(sorted(ALLOWED_SOURCE_PAGES))
31
+ return {"success": False, "message": f"source_page must be one of: {allowed}"}, 400
32
+ source_page = normalized_source_page
33
+
34
+ has_text = isinstance(text, str) and text.strip() != ""
35
+ has_vector = isinstance(vector, list) and len(vector) > 0
36
+ if not has_text and not has_vector:
37
+ return {"success": False, "message": "Missing required field: text or vector must be provided"}, 400
38
+
39
+ if token_index is not None and not isinstance(token_index, int):
40
+ return {"success": False, "message": "token_index must be an integer"}, 400
41
+ if token_index is not None and token_index < 0:
42
+ return {"success": False, "message": "token_index must be >= 0"}, 400
43
+ if has_text and token_index is None:
44
+ return {"success": False, "message": "token_index is required when text is provided"}, 400
45
+
46
+ client_ip = get_client_ip()
47
+ start_time = time.perf_counter()
48
+ request_id = log_prediction_attribute_request(
49
+ context=text if has_text else str(vector)[:200],
50
+ target_prediction=None,
51
+ target_token_id=token_index,
52
+ model=model,
53
+ source_page=source_page,
54
+ flow_id=None,
55
+ flow_step=None,
56
+ client_ip=client_ip,
57
+ )
58
+
59
+ lock_acquired = tiny_nla_lock.acquire(timeout=TINY_NLA_LOCK_TIMEOUT)
60
+ if not lock_acquired:
61
+ return {"success": False, "message": f"Tiny-NLA queue wait exceeded {TINY_NLA_LOCK_TIMEOUT} seconds; please try again later."}, 503
62
+
63
+ try:
64
+ engine = TinyNLAEngine()
65
+
66
+ if has_vector:
67
+ activation = torch.tensor(vector, dtype=torch.float32)
68
+ if activation.shape[0] != 1024:
69
+ return {"success": False, "message": f"vector dimension must be 1024, got {activation.shape[0]}"}, 400
70
+ else:
71
+ activation = engine.extract_activation(text, token_index)
72
+
73
+ explanation = engine.explain(activation)
74
+ roundtrip_cosine = engine.reconstruct_cosine(activation, explanation)
75
+
76
+ result = {
77
+ "concept": "",
78
+ "explanation": explanation,
79
+ "roundtrip_cosine": round(roundtrip_cosine, 4),
80
+ "vector_dim": 1024,
81
+ "note": "",
82
+ }
83
+ except ValueError as e:
84
+ return {"success": False, "message": str(e)}, 400
85
+ except Exception as e:
86
+ import traceback
87
+ traceback.print_exc()
88
+ exit_if_oom(e, defer_seconds=1)
89
+ return {"success": False, "message": str(e)}, 500
90
+ finally:
91
+ tiny_nla_lock.release()
92
+ gc.collect()
93
+
94
+ elapsed = time.perf_counter() - start_time
95
+ print(
96
+ f"\t📤 API activation_explain response: req_id={request_id}, "
97
+ f"concept={result.get('concept')!r}, roundtrip={result.get('roundtrip_cosine')}, "
98
+ f"response_time={elapsed:.4f}s"
99
+ )
100
+ return {"success": True, **result}, 200
backend/core/tiny_nla.py ADDED
@@ -0,0 +1,188 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """Tiny-NLA 引擎:LoRA + AR head 模型管理(lazy 单例)。
2
+
3
+ 职责:
4
+ - lazy 加载 Qwen/Qwen3-0.6B-Base(float32 + eager,对齐训练)+ LoRA adapter + AR head
5
+ - extract_activation(text, token_index) → layer 19 残差流
6
+ - explain(activation) → 注入激活 → generate → 自然语言解释
7
+ - reconstruct_cosine(activation, explanation) → AR head 重建 → cosine
8
+
9
+ 所有模型均独立加载(float32 + eager),不复用 base 槽(float16),确保与训练精度对齐。
10
+ """
11
+
12
+ import threading
13
+ import time
14
+ from pathlib import Path
15
+
16
+ import torch
17
+ import torch.nn.functional as F
18
+ import yaml
19
+ from peft import PeftModel
20
+ from transformers import AutoModelForCausalLM, AutoTokenizer
21
+
22
+ from backend.models.device import DeviceManager
23
+
24
+ tiny_nla_lock = threading.Lock()
25
+ TINY_NLA_LOCK_TIMEOUT = 30.0
26
+
27
+ REPO_ROOT = Path(__file__).resolve().parents[2]
28
+ SIDECAR_PATH = REPO_ROOT / "experiments" / "tiny_nla" / "nla_meta.yaml"
29
+ CHECKPOINT_DIR = REPO_ROOT / "artifacts" / "tiny_nla" / "checkpoints"
30
+ AV_CHECKPOINT = CHECKPOINT_DIR / "av"
31
+ AR_CHECKPOINT = CHECKPOINT_DIR / "ar" / "best_ar_head.pt"
32
+
33
+
34
+ class TinyNLAEngine:
35
+ """LoRA + AR head 单例引擎。首次调用时 lazy 加载全部模型。"""
36
+
37
+ _instance = None
38
+ _init_done = False
39
+
40
+ def __new__(cls):
41
+ if cls._instance is None:
42
+ cls._instance = super().__new__(cls)
43
+ return cls._instance
44
+
45
+ def __init__(self):
46
+ if self._init_done:
47
+ return
48
+ with open(SIDECAR_PATH, "r") as f:
49
+ self.meta = yaml.safe_load(f)
50
+
51
+ self.base_model_name = self.meta["base_model"]
52
+ self.av_model_name = self.meta["av_init_model"]
53
+ self.layer_idx = self.meta["layer_index"]
54
+ self.d_model = self.meta["d_model"]
55
+ self.inj_char = self.meta["tokens"]["injection_char"]
56
+ self.inj_token_id = self.meta["tokens"]["injection_token_id"]
57
+ self.inj_scale = self.meta["extraction"]["injection_scale"]
58
+
59
+ self.device = DeviceManager.get_device()
60
+ self.dtype = torch.float32
61
+
62
+ self._base_model = None
63
+ self._base_tokenizer = None
64
+ self._av_model = None
65
+ self._av_tokenizer = None
66
+ self._ar_head = None
67
+ self._init_done = True
68
+
69
+ def _ensure_loaded(self):
70
+ if self._base_model is not None:
71
+ return
72
+
73
+ t0 = time.perf_counter()
74
+ print(f" [TinyNLA] Loading base ({self.base_model_name})...")
75
+ self._base_model = AutoModelForCausalLM.from_pretrained(
76
+ self.base_model_name,
77
+ trust_remote_code=True,
78
+ torch_dtype=self.dtype,
79
+ low_cpu_mem_usage=True,
80
+ attn_implementation="eager",
81
+ ).to(self.device)
82
+ self._base_model.eval()
83
+ self._base_tokenizer = AutoTokenizer.from_pretrained(
84
+ self.base_model_name, trust_remote_code=True
85
+ )
86
+ DeviceManager.print_model_load_stats(self._base_model, time.perf_counter() - t0)
87
+
88
+ t1 = time.perf_counter()
89
+ print(f" [TinyNLA] Loading AV base ({self.av_model_name}) + LoRA...")
90
+ av_base = AutoModelForCausalLM.from_pretrained(
91
+ self.av_model_name,
92
+ trust_remote_code=True,
93
+ torch_dtype=self.dtype,
94
+ low_cpu_mem_usage=True,
95
+ attn_implementation="eager",
96
+ ).to(self.device)
97
+ self._av_model = PeftModel.from_pretrained(av_base, str(AV_CHECKPOINT))
98
+ self._av_model.eval()
99
+ self._av_tokenizer = AutoTokenizer.from_pretrained(
100
+ str(AV_CHECKPOINT), trust_remote_code=True
101
+ )
102
+ DeviceManager.print_model_load_stats(self._av_model, time.perf_counter() - t1)
103
+
104
+ if AR_CHECKPOINT.exists():
105
+ print(f" [TinyNLA] Loading AR head...")
106
+ state = torch.load(
107
+ str(AR_CHECKPOINT), map_location=self.device, weights_only=True
108
+ )
109
+ if "linear.weight" in state:
110
+ state = {"weight": state["linear.weight"]}
111
+ self._ar_head = torch.nn.Linear(self.d_model, self.d_model, bias=False)
112
+ self._ar_head.load_state_dict(state)
113
+ self._ar_head.to(self.device)
114
+ self._ar_head.eval()
115
+ else:
116
+ print(f" [TinyNLA] ⚠ AR checkpoint not found at {AR_CHECKPOINT}")
117
+ self._ar_head = None
118
+
119
+ def extract_activation(self, text: str, token_index: int) -> torch.Tensor:
120
+ """用 float32 base 模型提取 layer 19 残差流。"""
121
+ self._ensure_loaded()
122
+ inputs = self._base_tokenizer(text, return_tensors="pt").to(self.device)
123
+ seq_len = inputs["input_ids"].shape[1]
124
+ if token_index >= seq_len:
125
+ raise ValueError(f"token_index {token_index} out of range (seq_len={seq_len})")
126
+ with torch.no_grad():
127
+ outputs = self._base_model(**inputs, output_hidden_states=True, use_cache=False)
128
+ DeviceManager.synchronize(self.device)
129
+ activation = outputs.hidden_states[self.layer_idx][0, token_index, :].cpu()
130
+ return activation
131
+
132
+ def explain(self, activation: torch.Tensor, max_new_tokens: int = 64) -> str:
133
+ """注入激活向量 → generate → 返回 explanation 文本。"""
134
+ self._ensure_loaded()
135
+
136
+ prompt = f"<concept>{self.inj_char}</concept>\n<explanation>"
137
+ inputs = self._av_tokenizer(prompt, return_tensors="pt").to(self.device)
138
+
139
+ embeds = self._av_model.get_input_embeddings()(inputs["input_ids"])
140
+ inj_positions = (inputs["input_ids"][0] == self.inj_token_id).nonzero(as_tuple=True)[0]
141
+
142
+ if len(inj_positions) > 0:
143
+ inj_pos = inj_positions[0].item()
144
+ norm = activation.norm()
145
+ if norm > 0:
146
+ activation = activation / norm * self.inj_scale
147
+ scaled_act = activation.to(embeds.dtype).to(self.device)
148
+ embeds[0, inj_pos, :] = scaled_act
149
+
150
+ with torch.no_grad():
151
+ output_ids = self._av_model.generate(
152
+ inputs_embeds=embeds,
153
+ max_new_tokens=max_new_tokens,
154
+ do_sample=False,
155
+ pad_token_id=self._av_tokenizer.pad_token_id or self._av_tokenizer.eos_token_id,
156
+ )
157
+
158
+ DeviceManager.synchronize(self.device)
159
+
160
+ prompt_len_tokens = inputs["input_ids"].shape[1]
161
+ gen_token_ids = output_ids[0][prompt_len_tokens:]
162
+ explanation = self._av_tokenizer.decode(gen_token_ids, skip_special_tokens=True).strip()
163
+ return explanation
164
+
165
+ def reconstruct_cosine(self, activation: torch.Tensor, explanation: str) -> float:
166
+ """AR head 重建 → 计算 cosine。用 float32 base 模型提取 explanation 的 last hidden。"""
167
+ self._ensure_loaded()
168
+ if self._ar_head is None:
169
+ return 0.0
170
+
171
+ inputs = self._base_tokenizer(
172
+ explanation, return_tensors="pt", truncation=True, max_length=128
173
+ ).to(self.device)
174
+
175
+ with torch.no_grad():
176
+ outputs = self._base_model(**inputs, output_hidden_states=True, use_cache=False)
177
+ last_hidden = outputs.hidden_states[-1]
178
+ seq_len = inputs["attention_mask"].sum(dim=1) - 1
179
+ last_token_hidden = last_hidden[0, seq_len[0], :]
180
+
181
+ reconstructed = self._ar_head(last_token_hidden)
182
+
183
+ DeviceManager.synchronize(self.device)
184
+
185
+ orig_n = F.normalize(activation.unsqueeze(0).to(self.device), dim=-1)
186
+ recon_n = F.normalize(reconstructed.unsqueeze(0), dim=-1)
187
+ cosine = (orig_n * recon_n).sum(dim=-1).item()
188
+ return cosine
backend/platform/source_page.py CHANGED
@@ -1,6 +1,6 @@
1
  """prediction_attribute API 的 source_page slug 规范化(含旧客户端兼容)。"""
2
 
3
- ALLOWED_SOURCE_PAGES = frozenset({"analysis", "chat", "attribution", "causal_flow"})
4
 
5
 
6
  def normalize_source_page(raw: str) -> str | None:
 
1
  """prediction_attribute API 的 source_page slug 规范化(含旧客户端兼容)。"""
2
 
3
+ ALLOWED_SOURCE_PAGES = frozenset({"analysis", "chat", "attribution", "causal_flow", "logit_lens"})
4
 
5
 
6
  def normalize_source_page(raw: str) -> str | None:
client/src/css/pages/logit_lens.scss CHANGED
@@ -84,6 +84,54 @@
84
  margin-bottom: 4px;
85
  }
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
87
  // ---- 层卡片 ----
88
  .ll-layer-card {
89
  margin-top: 12px;
 
84
  margin-bottom: 4px;
85
  }
86
 
87
+ // ---- Activation Explainer ----
88
+ .ae-loading {
89
+ font-size: 9pt;
90
+ color: var(--text-muted);
91
+ padding: 12px 0;
92
+ text-align: center;
93
+ }
94
+
95
+ .ae-result {
96
+ margin-top: 8px;
97
+ }
98
+
99
+ .ae-field {
100
+ margin-bottom: 8px;
101
+ display: flex;
102
+ align-items: flex-start;
103
+ gap: 8px;
104
+ }
105
+
106
+ .ae-label {
107
+ font-size: 10pt;
108
+ font-weight: 600;
109
+ color: var(--text-secondary);
110
+ min-width: 48px;
111
+ flex-shrink: 0;
112
+ }
113
+
114
+ .ae-value {
115
+ font-size: 10pt;
116
+ color: var(--text-primary);
117
+ line-height: 1.5;
118
+ }
119
+
120
+ .ae-cosine {
121
+ font-weight: 700;
122
+ font-size: 11pt;
123
+ &.cosine-excellent { color: #16a34a; }
124
+ &.cosine-good { color: #2563eb; }
125
+ &.cosine-fair { color: #d97706; }
126
+ &.cosine-low { color: #dc2626; }
127
+ }
128
+
129
+ .ae-error {
130
+ font-size: 9pt;
131
+ color: #dc2626;
132
+ padding: 8px 0;
133
+ }
134
+
135
  // ---- 层卡片 ----
136
  .ll-layer-card {
137
  margin-top: 12px;
client/src/logit_lens.html CHANGED
@@ -84,6 +84,16 @@
84
  </div>
85
  </section>
86
  </div>
 
 
 
 
 
 
 
 
 
 
87
  </section>
88
 
89
  <div class="resizer" id="resizer"></div>
 
84
  </div>
85
  </section>
86
  </div>
87
+
88
+ <div class="logit-lens-panel" id="ae_panel" style="display:none;">
89
+ <div class="attribution-panel-label">Activation Explainer / 激活解释</div>
90
+ <div id="ae_loading" class="ae-loading" style="display:none">🧠 正在推理... 激活解释需要较长时间(约 5-15 秒)</div>
91
+ <div id="ae_result" class="ae-result" style="display:none">
92
+ <div class="ae-field"><span class="ae-label">解释</span><span id="ae_explanation" class="ae-value"></span></div>
93
+ <div class="ae-field"><span class="ae-label">可信度</span><span id="ae_cosine" class="ae-cosine"></span><span>&nbsp;(roundtrip cosine)</span></div>
94
+ </div>
95
+ <div id="ae_error" class="ae-error" style="display:none"></div>
96
+ </div>
97
  </section>
98
 
99
  <div class="resizer" id="resizer"></div>
client/src/pages/logit_lens/index.ts CHANGED
@@ -16,7 +16,7 @@ import { showAlertDialog } from '../../shared/ui/dialog';
16
  import URLHandler from '../../shared/core/URLHandler';
17
  import { createToast } from '../../shared/ui/toast';
18
  import { translateApiErrorMessage } from '../../shared/core/errorUtils';
19
- import type { LogitLensResult, LogitLensLayer } from '../../shared/api/GLTR_API';
20
  import { lsReadEnum, lsWriteString } from '../../shared/storage/localStorageHelpers';
21
 
22
  d3.selectAll('.loadersmall').style('display', 'none');
@@ -355,6 +355,61 @@ function renderLayerCard(): void {
355
  </div>`;
356
  }
357
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
358
  // --- 主分析逻辑 ---
359
  async function runAnalyze(): Promise<void> {
360
  const context = (contextField.node() as HTMLTextAreaElement | null)?.value ?? '';
@@ -369,6 +424,11 @@ async function runAnalyze(): Promise<void> {
369
  lastCommittedInputs = { context, target };
370
  const info = `${tr('model')}: ${ll.model ?? '–'}\n${tr('Target token')}: ${ll.target_token ?? '–'}\n${tr('Final prob')}: ${(ll.final_target_prob * 100).toFixed(1)}%\n${tr('Layers')}: ${ll.n_layers ?? '–'}`;
371
  resultInfoEl.classed('is-hidden', false).text(info);
 
 
 
 
 
372
  } else {
373
  showAlertDialog(tr('Logit Lens'), ll.message || tr('Analysis failed'));
374
  }
 
16
  import URLHandler from '../../shared/core/URLHandler';
17
  import { createToast } from '../../shared/ui/toast';
18
  import { translateApiErrorMessage } from '../../shared/core/errorUtils';
19
+ import type { LogitLensResult, LogitLensLayer, ActivationExplainResult } from '../../shared/api/GLTR_API';
20
  import { lsReadEnum, lsWriteString } from '../../shared/storage/localStorageHelpers';
21
 
22
  d3.selectAll('.loadersmall').style('display', 'none');
 
355
  </div>`;
356
  }
357
 
358
+ // --- Activation Explainer ---
359
+ async function runActivationExplain(context: string, targetToken: string): Promise<void> {
360
+ const panel = document.getElementById('ae_panel');
361
+ const loading = document.getElementById('ae_loading');
362
+ const result = document.getElementById('ae_result');
363
+ const errEl = document.getElementById('ae_error');
364
+ if (!panel || !loading || !result || !errEl) return;
365
+
366
+ panel.style.display = 'block';
367
+ loading!.style.display = '';
368
+ result.style.display = 'none';
369
+ errEl.style.display = 'none';
370
+
371
+ try {
372
+ // 1) tokenize 找出 target token 的 index
373
+ const tok = await api.tokenize(context, currentModelVariant());
374
+ const spans = tok?.spans ?? [];
375
+ let tokenIndex = -1;
376
+ for (let i = 0; i < spans.length; i++) {
377
+ if (spans[i].raw === targetToken) { tokenIndex = i; break; }
378
+ }
379
+ if (tokenIndex < 0) {
380
+ // fallback:取最后一个 token
381
+ tokenIndex = spans.length - 1;
382
+ }
383
+
384
+ // 2) call activation-explain
385
+ const ae = await api.explainActivation(currentModelVariant(), 'logit_lens', context, tokenIndex);
386
+ if (!ae.success) {
387
+ errEl.textContent = ae.message || tr('Activation explain failed');
388
+ errEl.style.display = '';
389
+ return;
390
+ }
391
+
392
+ // 3) render
393
+ const cosine = ae.roundtrip_cosine ?? 0;
394
+ let cls = 'cosine-low';
395
+ if (cosine >= 0.70) cls = 'cosine-excellent';
396
+ else if (cosine >= 0.60) cls = 'cosine-good';
397
+ else if (cosine >= 0.50) cls = 'cosine-fair';
398
+ const cosineHtml = `<span class="ae-cosine ${cls}">${(cosine * 100).toFixed(1)}%</span>`;
399
+
400
+ const explanationEl = document.getElementById('ae_explanation');
401
+ const cosineEl = document.getElementById('ae_cosine');
402
+ if (explanationEl) explanationEl.textContent = ae.explanation ?? '';
403
+ if (cosineEl) cosineEl.innerHTML = cosineHtml;
404
+ result.style.display = '';
405
+ } catch (err: unknown) {
406
+ errEl.textContent = (err instanceof Error ? err.message : String(err));
407
+ errEl.style.display = '';
408
+ } finally {
409
+ loading!.style.display = 'none';
410
+ }
411
+ }
412
+
413
  // --- 主分析逻辑 ---
414
  async function runAnalyze(): Promise<void> {
415
  const context = (contextField.node() as HTMLTextAreaElement | null)?.value ?? '';
 
424
  lastCommittedInputs = { context, target };
425
  const info = `${tr('model')}: ${ll.model ?? '–'}\n${tr('Target token')}: ${ll.target_token ?? '–'}\n${tr('Final prob')}: ${(ll.final_target_prob * 100).toFixed(1)}%\n${tr('Layers')}: ${ll.n_layers ?? '–'}`;
426
  resultInfoEl.classed('is-hidden', false).text(info);
427
+
428
+ // 自动触发 Activation Explainer
429
+ if (context && ll.target_token) {
430
+ void runActivationExplain(context, ll.target_token);
431
+ }
432
  } else {
433
  showAlertDialog(tr('Logit Lens'), ll.message || tr('Analysis failed'));
434
  }
client/src/shared/api/GLTR_API.ts CHANGED
@@ -102,6 +102,16 @@ export type BranchNextResult = {
102
  message?: string;
103
  };
104
 
 
 
 
 
 
 
 
 
 
 
105
  export class TextAnalysisAPI {
106
  private adminToken: string | null = null;
107
 
@@ -504,6 +514,51 @@ export class TextAnalysisAPI {
504
  return data as BranchNextResult;
505
  }
506
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
507
  /**
508
  * 使用SSE流式分析文本,支持进度回调(内部方法)
509
  */
 
102
  message?: string;
103
  };
104
 
105
+ export type ActivationExplainResult = {
106
+ success: boolean;
107
+ concept?: string;
108
+ explanation?: string;
109
+ roundtrip_cosine?: number;
110
+ vector_dim?: number;
111
+ note?: string;
112
+ message?: string;
113
+ };
114
+
115
  export class TextAnalysisAPI {
116
  private adminToken: string | null = null;
117
 
 
514
  return data as BranchNextResult;
515
  }
516
 
517
+ /**
518
+ * Tokenize:将文本分词为 token spans。
519
+ */
520
+ public async tokenize(
521
+ context: string,
522
+ model: string,
523
+ signal?: AbortSignal
524
+ ): Promise<{ success: boolean; spans: Array<{ offset: [number, number]; raw: string; token_id?: number }> }> {
525
+ const res = await fetch(this.baseURL + '/api/tokenize', {
526
+ method: 'POST',
527
+ headers: this.getHeaders(),
528
+ body: JSON.stringify({ context, model }),
529
+ signal
530
+ });
531
+ const data = await res.json();
532
+ if (data && data.success === false) {
533
+ throw new Error(data.message || 'Tokenize failed');
534
+ }
535
+ return data;
536
+ }
537
+
538
+ /**
539
+ * Activation Explainer (Tiny-NLA):解释激活向量的概念语义。
540
+ */
541
+ public async explainActivation(
542
+ model: string,
543
+ sourcePage: string,
544
+ text: string,
545
+ tokenIndex: number,
546
+ signal?: AbortSignal
547
+ ): Promise<ActivationExplainResult> {
548
+ const bodyObj: Record<string, unknown> = { model, source_page: sourcePage, text, token_index: tokenIndex };
549
+ const res = await fetch(this.baseURL + '/api/activation-explain', {
550
+ method: 'POST',
551
+ headers: this.getHeaders(),
552
+ body: JSON.stringify(bodyObj),
553
+ signal
554
+ });
555
+ const data = await res.json();
556
+ if (data && data.success === false) {
557
+ throw new Error(data.message || 'Activation explain failed');
558
+ }
559
+ return data as ActivationExplainResult;
560
+ }
561
+
562
  /**
563
  * 使用SSE流式分析文本,支持进度回调(内部方法)
564
  */
client/src/shared/lang/translations.ts CHANGED
@@ -512,5 +512,10 @@ export const translations: Translations = {
512
  'which input tokens drove this prediction? Gradient scores each context token by how much changing it would shift the output probability.': '哪些输入词推动了这个预测?梯度法通过计算每个 token 对输出概率的扰动敏感度,量化其贡献。',
513
  'applies the final output layer to each intermediate layer, revealing how the model\'s prediction takes shape across depth — early layers guess, later layers converge.': '将最终输出层应用于每个中间层,揭示预测如何随深度逐步成型——浅层在猜测,深层在收敛。',
514
  'shows the model\'s top-k next-token candidates at each step as an interactive tree — revealing how probability mass spreads across possible continuations.': '将每一步的 top-k 候选词展开为可交互树,揭示概率质量如何在可能的续写路径间分布。',
 
 
 
 
 
515
  }
516
  };
 
512
  'which input tokens drove this prediction? Gradient scores each context token by how much changing it would shift the output probability.': '哪些输入词推动了这个预测?梯度法通过计算每个 token 对输出概率的扰动敏感度,量化其贡献。',
513
  'applies the final output layer to each intermediate layer, revealing how the model\'s prediction takes shape across depth — early layers guess, later layers converge.': '将最终输出层应用于每个中间层,揭示预测如何随深度逐步成型——浅层在猜测,深层在收敛。',
514
  'shows the model\'s top-k next-token candidates at each step as an interactive tree — revealing how probability mass spreads across possible continuations.': '将每一步的 top-k 候选词展开为可交互树,揭示概率质量如何在可能的续写路径间分布。',
515
+
516
+ // ---- Activation Explainer ----
517
+ 'Activation Explainer': '激活解释',
518
+ 'Activation explain failed': '激活解释失败',
519
+ 'Loading activation explanation…': '正在推理激活解释…',
520
  }
521
  };
experiments/tiny_nla/PLAN_v2.md ADDED
@@ -0,0 +1,406 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Tiny NLA v2 Implementation Plan
2
+
3
+ > 目标:修复三个关键问题,让 Tiny NLA 项目从「SFT 概念验证」升级到「符合 Anthropic NLA 方案的完整训练闭环」
4
+
5
+ ## 问题总览
6
+
7
+ | 优先级 | 问题 | 当前状态 | 目标状态 |
8
+ |:------:|------|----------|----------|
9
+ | P0 | 缺少 RL(GRPO)训练循环 | 只有 SFT,没有 RL | 实现 AV+AR 联合 GRPO RL,reward = `-mse_nrm` |
10
+ | P0 | Teacher label 带了太多辅助信息 | prompt 包含 context、top_tokens、pos 等 | 简化到仅依赖激活向量本身的信息 |
11
+ | P1 | 数据量太小(1461 vs 1M) | 22 条句子,284→1461 records | 扩大到 ≥10K 条 records |
12
+
13
+ ---
14
+
15
+ ## Phase 0: 数据扩容(P1 — 数据量)
16
+
17
+ > **依赖**:无,可独立进行
18
+ > **预计耗时**:4-6 小时(大部分是 API 调用等待时间)
19
+
20
+ ### Step 0.1: 扩大语料库到 ≥500 条句子
21
+
22
+ **文件**:`experiments/tiny_nla/expand_dataset_v2.py`(新文件)
23
+
24
+ 当前 `expand_dataset.py` 用 15 个领域生成 ~300 条句子。需要扩大到 500+ 条。
25
+
26
+ **策略**:
27
+ - 扩展领域列表到 30+ 个(增加:心理学、心理学、编程技术、数学逻辑、影视娱乐、宗教信仰、语言语言学、艺术美学、交通出行、宠物动物、天气气候、金融投资、职场沟通)
28
+ - 每个领域生成 20-30 条
29
+ - 用 opencode CLI 并发批量生成(参考现有 `batch_teacher_labels.py` 的并发模式)
30
+ - 增加句子多样性:短句(8-15字)、中等句(15-30字)、长句(30-50字)
31
+
32
+ **输出**:`artifacts/tiny_nla/expanded_texts_v2.json`(500+ 条句子)
33
+
34
+ ### Step 0.2: 提取全部激活向量
35
+
36
+ **文件**:`experiments/tiny_nla/generate_data_v2.py`(新文件,基于现有 `generate_data.py`)
37
+
38
+ - 加载 Qwen3-0.6B-Base
39
+ - 对 500+ 条句子做 forward pass,提取 layer 19 的残差流
40
+ - 每条句子每个 token 都提取(不需要过滤)
41
+ - 保存为 parquet 格式(与 Anthropic NLA 对齐)
42
+
43
+ **输出**:`artifacts/tiny_nla/activations_v2.parquet`
44
+
45
+ **预估**:500 条句子 × ~15 tokens/句 = ~7,500 activation records。以 0.6B 模型在 M4 Pro 上,forward pass 极快(<5 分钟全部完成)。
46
+
47
+ ### Step 0.3: 过滤 OOD 样本
48
+
49
+ - 去掉 activation_norm > 2000 的首 token(与现有做法一致)
50
+ - 去掉 norm < 50 的异常低激活(可能是 padding/特殊 token)
51
+ - 预估保留 ~7,000 条 in-distribution records
52
+
53
+ ---
54
+
55
+ ## Phase 1: Teacher Label 重做(P0 — 辅助信息)
56
+
57
+ > **依赖**:Phase 0 完成后的数据
58
+ > **预计耗时**:8-12 小时(API 调用)
59
+
60
+ ### Step 1.1: 设计简化版 Teacher Prompt
61
+
62
+ **核心原则**:NLA 的目标是让 AV **仅从激活向量**推断语义。Teacher label 的训练数据不应包含 AV 运行时看不到的信息。
63
+
64
+ **Anthropic 方案对比**:
65
+ Anthropic 的 teacher 生成用的是 **Claude 看 activation 向量的投影/统计信息**,不附带原始句子或 top_tokens。
66
+
67
+ **新 prompt 设计**:
68
+
69
+ ```
70
+ 你是一个语言模型内部机制分析专家。我给你一段来自 Qwen3-0.6B 模型(layer 19,residual stream)
71
+ 在某 token 位置的激活向量信息,请你推断这个位置编码了什么语义。
72
+
73
+ 激活向量统计信息:
74
+ - L2 norm: {norm}
75
+ - 与平均激活方向的余弦相似度: {cosine_with_mean}
76
+ - 该位置在序列中的相对位置: {relative_pos}(0=开头, 1=末尾)
77
+
78
+ 请用1-2句简洁中文描述你认为这个位置编码的语义信息。
79
+ 要求:简洁具体,不超过55字,只输出解释本身。
80
+ ```
81
+
82
+ **关键变化**:
83
+ - ❌ 去掉原始句子文本(context)
84
+ - ❌ 去掉 top_tokens(模型预测候选)
85
+ - ❌ 去掉 token_text(当前 token 文本)
86
+ - ✅ 只给激活向量的**统计特征**(norm、方向、相对位置)
87
+ - ✅ 这些信息 AV 在推理时可以通过注入的向量获得
88
+
89
+ **为什么可以去掉 context/top_tokens**:
90
+ - Anthropic 的论文明确指出,NLA 的核心价值在于**从向量本身推断语义**
91
+ - 如果 teacher 依赖了上下文,AV 就不需要学会「读懂向量」,而是学会了「从 prompt 里偷看上下文」
92
+ - RL 训练时,reward 只看「向量→文本→向量」的 roundtrip 质量,上下文信息不会被注入
93
+
94
+ ### Step 1.2: 批量生成简化版 Teacher Labels
95
+
96
+ **文件**:`experiments/tiny_nla/generate_teacher_labels_v2.py`(新文件)
97
+
98
+ **API 方案:双通道并发**
99
+
100
+ 为最大化吞吐量,同时使用两个免费 API 通道:
101
+
102
+ | 通道 | API | 模型 | 并发数 | 特点 |
103
+ |------|-----|------|:------:|------|
104
+ | 通道 A | opencode CLI | deepseek-v4-flash-free | 4 workers | 本地 CLI 调用,已在用 |
105
+ | 通道 B | NVIDIA API | deepseek-ai/deepseek-v4-pro | 8-10 concurrent | 直接 HTTP API,更快更强 |
106
+
107
+ **NVIDIA API 配置**:
108
+ ```python
109
+ from openai import OpenAI
110
+
111
+ nvidia_client = OpenAI(
112
+ base_url="https://integrate.api.nvidia.com/v1",
113
+ api_key=os.environ.get("NVIDIA_API_KEY", ""),
114
+ )
115
+ # model: "deepseek-ai/deepseek-v4-pro"
116
+ # thinking=False → 直接输出,不需要等思考过程
117
+ # max_tokens=256 足够(teacher label ≤55字)
118
+ ```
119
+
120
+ **并发策略**:
121
+ - opencode CLI:4 个 subprocess worker(需要达到 free tier 并发上限,我不确定4个是否是上限)
122
+ - NVIDIA API:8-10 个 asyncio concurrent request(NVIDIA free tier 通常允许较高并发)
123
+ - **总计 12-14 个并发请求**,吞吐量提升约 3 倍
124
+
125
+ **优势对比**:
126
+ | 方案 | 速度 | 模型质量 | 成本 |
127
+ |------|------|---------|------|
128
+ | opencode CLI (flash-free) | ~3-5 秒/条 | DeepSeek V4 Flash | 免费 |
129
+ | NVIDIA API (v4-pro) | ~2-3 秒/条 | DeepSeek V4 Pro(更强) | 免费 |
130
+ | **双通道合计** | **~12-14 条/分钟** | 混合质量 | **免费** |
131
+
132
+ **容错与 checkpoint**:
133
+ - 每 100 条保存 checkpoint(JSON 文件)
134
+ - 两个通道的结果合并写入同一个文件
135
+ - 支持 `--skip-existing`,中断后重启自动跳过已完成
136
+ - NVIDIA API 失败时自动重试 3 次,最终 fallback 到 opencode
137
+
138
+ **预估**:7000 条 ÷ 20 条/分钟 = **~6 小时**(比纯 opencode 方案快 ~5 倍)
139
+
140
+ **输出**:`artifacts/tiny_nla/teacher_labels_v2.json`
141
+
142
+ ### Step 1.3: 质量检查
143
+
144
+ - 检查是否有空输出或幻觉(声称上下文包含某些不存在的内容)
145
+ - 剔除长度 > 55 字的过长解释
146
+ - 剔除明显泛泛而谈的解释(如「编码了关键语义信息」「是核心语义单元」等套话)
147
+ - 统计剩余有效 label 数量
148
+
149
+ ---
150
+
151
+ ## Phase 2: SFT 重训(基于新数据)
152
+
153
+ > **依赖**:Phase 0 + Phase 1
154
+ > **预计耗时**:30-40 分钟
155
+
156
+ ### Step 2.1: AR SFT 重训
157
+
158
+ **文件**:修改现有 `train_ar.py`
159
+
160
+ - 使用新数据(Phase 0 的 7K+ records + Phase 1 的 teacher labels)
161
+ - AR 输入:teacher explanation text
162
+ - AR 输出:L2-normalized activation vector
163
+ - Loss: MSE on L2-normalized vectors
164
+ - 增加 train/val split (90/10)
165
+ - 训练 5-10 epochs,early stopping on val loss
166
+
167
+ **基线**:
168
+ - Mean baseline(永远预测均值方向)
169
+ - Shuffled baseline(随机配对 explanation→activation)
170
+
171
+ ### Step 2.2: AV SFT 重训
172
+
173
+ **文件**:修改现有 `train_av.py`
174
+
175
+ - 使用新数据
176
+ - AV 输入:injected activation vector (injection token ㈎)
177
+ - AV 输出:teacher explanation text
178
+ - LoRA 微调 Qwen3-0.6B Instruct
179
+ - 增加数据量到 7K+ 条,训练更长
180
+
181
+ **验证**:AV→AR roundtrip cosine 应 > 0.6(Anthropic 的 7B 模型达到 ~0.75)
182
+
183
+ ---
184
+
185
+ ## Phase 3: GRPO RL 训练(P0 — RL 循环)
186
+
187
+ > **依赖**:Phase 2 完成后的 SFT checkpoints
188
+ > **预计耗时**:8-12 小时
189
+ > **这是最关键的新增部分**
190
+
191
+ ### Step 3.1: 设计 GRPO RL 训练框架
192
+
193
+ **文件**:`experiments/tiny_nla/train_rl_grpo.py`(新文件)
194
+
195
+ **GRPO(Group Relative Policy Optimization)核心逻辑**:
196
+
197
+ ```
198
+ 对于每个 training batch:
199
+ 1. 从数据集取一批 activation vectors
200
+ 2. 用 AV 对每个 vector 生成 K 个候选解释(K=4,不同采样策略)
201
+ 3. 用 AR 对每个候选解释还原向量
202
+ 4. 计算 reward = -mse_nrm(reconstructed, original)
203
+ 其中 mse_nrm = MSE(L2_norm(original), L2_norm(reconstructed))
204
+ 5. 在 K 个候选中,用相对排名做 advantage:
205
+ advantage_i = (reward_i - mean(rewards_group)) / std(rewards_group)
206
+ 6. 用 advantage 加权 AV 的 policy gradient loss
207
+ 7. 同时 AR 继续做 supervised MSE loss(保持 AR 能力不退化)
208
+ ```
209
+
210
+ ### Step 3.2: 内存管理策略(适配 M4 Pro 48GB)
211
+
212
+ **策略:逐模型加载(sequential rollout)**
213
+
214
+ ```python
215
+ class RLTrainer:
216
+ def __init__(self):
217
+ # 不一次性加载所有模型
218
+ # 按需加载/卸载
219
+
220
+ def rollout_phase(self, activations):
221
+ """Step 1: 加载 Target + AV,生成候选解释"""
222
+ self._load_target_and_av() # ~2.4GB
223
+ candidates = []
224
+ for act in activations:
225
+ for k in range(GROUP_SIZE):
226
+ explanation = self.av.generate(act, temperature=0.7 + k*0.1)
227
+ candidates.append((act, explanation))
228
+ self._unload_av()
229
+ return candidates
230
+
231
+ def reward_phase(self, candidates):
232
+ """Step 2: 加载 AR,计算 reward"""
233
+ self._load_ar() # ~1.2GB
234
+ rewards = []
235
+ for act, explanation in candidates:
236
+ reconstructed = self.ar.reconstruct(explanation)
237
+ mse = mse_nrm(act, reconstructed)
238
+ rewards.append(-mse)
239
+ self._unload_ar()
240
+ return rewards
241
+
242
+ def update_phase(self, candidates, rewards, advantages):
243
+ """Step 3: 加载 AV,做 GRPO policy gradient update"""
244
+ self._load_av_with_grad() # ~2.4GB + optimizer states
245
+ self.av.grpo_update(candidates, advantages)
246
+ self._unload_av()
247
+ ```
248
+
249
+ **内存峰值**:每次只加载 1-2 个模型,峰值 ~8-12GB,远低于 48GB 限制。
250
+
251
+ ### Step 3.3: MPS 适配细节
252
+
253
+ - 使用 `float32`(MPS 不完全支持 bf16 训练)
254
+ - 遇到 MPS 不支持的 op 时,fallback 到 CPU
255
+ - `torch.no_grad()` 包裹所有 inference phase
256
+ - 定期 `torch.mps.empty_cache()` ���理缓存
257
+ - 生成时用 `do_sample=True, temperature=0.7` 产生多样性
258
+
259
+ ### Step 3.4: 训练超参数
260
+
261
+ ```yaml
262
+ # RL 训练配置
263
+ group_size: 4 # 每个 activation 生成 4 个候选解释
264
+ batch_size: 8 # 每批 8 个 activations
265
+ learning_rate: 1e-5 # AV LoRA 学习率(比 SFT 小)
266
+ ar_learning_rate: 1e-4 # AR head 学习率
267
+ num_steps: 1000 # RL 训练步数
268
+ max_new_tokens: 64 # AV 生成最大 token 数
269
+ temperature_range: [0.7, 0.8, 0.9, 1.0] # K 个候选的温度
270
+ eval_interval: 50 # 每 50 步评估一次
271
+ save_interval: 100 # 每 100 步保存 checkpoint
272
+ ```
273
+
274
+ ### Step 3.5: 评估指标
275
+
276
+ 每 `eval_interval` 步,计算:
277
+ - **Roundtrip cosine**:AV→AR 的还原质量(越高越好)
278
+ - **Roundtrip MSE**:越小越好
279
+ - **Explanation quality**:解释的平均长度、多样性
280
+ - **Baselines**:
281
+ - Mean direction baseline cosine
282
+ - SFT-only baseline(不经过 RL 的 AV+AR cosine)
283
+ - Shuffled baseline
284
+
285
+ ### Step 3.6: 训练流程图
286
+
287
+ ```mermaid
288
+ flowchart TD
289
+ A["Phase 2: SFT Checkpoints"] --> B["加载 Target Model<br/>Qwen3-0.6B-Base frozen"]
290
+ B --> C["Rollout: AV 生成 K=4 候选解释"]
291
+ C --> D["Reward: AR 还原向量<br/>计算 -mse_nrm"]
292
+ D --> E["Advantage: 组内相对排名"]
293
+ E --> F["Update: AV GRPO policy gradient"]
294
+ F --> G["AR: supervised MSE update"]
295
+ G --> H{"每 50 步评估"}
296
+ H -->|"roundtrip cosine 提升"| I["继续训练"]
297
+ H -->|"收敛或退化"| J["Early stop"]
298
+ I --> C
299
+ J --> K["保存最佳 checkpoint"]
300
+ ```
301
+
302
+ ---
303
+
304
+ ## Phase 4: 最终评估与验证
305
+
306
+ > **依赖**:Phase 3 完成
307
+ > **预计耗时**:30 分钟
308
+
309
+ ### Step 4.1: 完整 Roundtrip 评估
310
+
311
+ **文件**:修改现有 `eval_roundtrip.py`
312
+
313
+ - 用最终 RL checkpoint 做 roundtrip 评估
314
+ - 对比 SFT-only 和 RL 版本
315
+ - 计算完整指标表:
316
+
317
+ | 指标 | SFT-only | RL (GRPO) | 目标 |
318
+ |------|----------|-----------|------|
319
+ | Roundtrip cosine | 0.60 (current) | > 0.65 | > 0.70 |
320
+ | Roundtrip MSE | 0.80 | < 0.60 | < 0.50 |
321
+ | Mean baseline cosine | 0.60 | - | - |
322
+ | Shuffled baseline cosine | 0.53 | - | - |
323
+
324
+ ### Step 4.2: 人类评估
325
+
326
+ - 取 20 个随机 activation,展示 AV 解释
327
+ - 人工判断解释是否合理、是否与该位置的语义相关
328
+ - 记录幻觉率(解释中声称上下文包含某些不存在的内容的比例)
329
+
330
+ ### Step 4.3: 更新 Demo
331
+
332
+ - 用 RL 版本的 AV 模型更新 demo.html
333
+ - 对比展示 SFT vs RL 的解释质量差异
334
+
335
+ ---
336
+
337
+ ## 文件清单(新增/修改)
338
+
339
+ | 文件 | 状态 | 说明 |
340
+ |------|------|------|
341
+ | `expand_dataset_v2.py` | 新增 | Phase 0.1:扩大语料库 |
342
+ | `generate_data_v2.py` | 新增 | Phase 0.2:提取激活向量(parquet) |
343
+ | `generate_teacher_labels_v2.py` | 新增 | Phase 1.2:简化版 teacher label |
344
+ | `train_ar.py` | 修改 | Phase 2.1:适配新数据格式 |
345
+ | `train_av.py` | 修改 | Phase 2.2:适配新数据格式 |
346
+ | `train_rl_grpo.py` | 新增 | Phase 3:GRPO RL 训练 |
347
+ | `eval_roundtrip.py` | 修改 | Phase 4:完整评估 |
348
+ | `nla_meta.yaml` | 修改 | 更新训练配置记录 |
349
+
350
+ ---
351
+
352
+ ## 时间估算总览
353
+
354
+ | Phase | 内容 | 预计耗时 |
355
+ |:------:|------|----------|
356
+ | 0 | 数据扩容 | 4-6 小时 |
357
+ | 1 | Teacher label 重做(三通道 API) | **~6 小时** |
358
+ | 2 | SFT 重训 | 30-40 分钟 |
359
+ | 3 | GRPO RL | 8-12 小时 |
360
+ | 4 | 评估 | 30 分钟 |
361
+ | **总计** | | **~19-25 小时** |
362
+
363
+ 大部分时间是 API 调用等待(Phase 1)和 RL 训练迭代(Phase 3),代码编写本身只需 2-3 小时。
364
+
365
+ ---
366
+
367
+ ## 风险与缓解
368
+
369
+ | 风险 | 影响 | 缓解方案 |
370
+ |------|------|---------|
371
+ | MPS 不支持某些 op | RL 训练 crash | fallback CPU + 测试 smoke case |
372
+ | 0.6B 模型太小,roundtrip cosine 无法提升到 >0.65 | RL 收益有限 | 先跑一轮看 baseline gap;如果 SFT cosine 已经很低则 RL 空间有限 |
373
+ | 简化版 teacher label 质量差(没有上下文辅助) | AV SFT 起点太低 | 可保留一个"半简化"版本作为对照:给 norm + top_tokens 但不给原文 |
374
+ | NVIDIA API rate limit / key 过期 | 生成速度降低 | fallback 到其他通道;API key 安全存储在 .env |
375
+ | OpenRouter free tier rate limit | nex-n2-pro:free 并发受限 | 降低并发数;fallback 到 NVIDIA/opencode |
376
+ | opencode CLI 并发上限 | 无法同时调太多 | 已达到 4 workers 上限,NVIDIA + OpenRouter 补充并发 |
377
+
378
+ ---
379
+
380
+ ## API Key 安全策略
381
+
382
+ 所有 API key 存储在 `.env` 文件中(不硬编码到代码):
383
+
384
+ ```bash
385
+ # InfoLens/.env 文件(不纳入 git,已在 .gitignore 中)
386
+ NVIDIA_API_KEY=<your-key-here>
387
+ OPENROUTER_API_KEY=<your-key-here>
388
+ ```
389
+
390
+ 代码中通过 `os.environ.get()` 或 `python-dotenv` 读取:
391
+ ```python
392
+ from dotenv import load_dotenv
393
+ load_dotenv() # 自动加载 .env
394
+ nvidia_key = os.environ.get("NVIDIA_API_KEY")
395
+ openrouter_key = os.environ.get("OPENROUTER_API_KEY")
396
+ ```
397
+
398
+ ---
399
+
400
+ ## API Key 安全策略
401
+
402
+ NVIDIA API key 不应硬编码在代码中,而是:
403
+ - 存储在 `.env` 文件中:`NVIDIA_API_KEY=nvapi-...`
404
+ - 代码中通过 `os.environ.get("NVIDIA_API_KEY")` 读取
405
+ - `.env` 文件不纳入 git(已在 `.gitignore` 中)
406
+ - Plan 文档中只记录使用方式,不记录完整 key |
experiments/tiny_nla/batch_teacher_labels.py ADDED
@@ -0,0 +1,105 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Batch teacher label generation via opencode (parallel workers).
4
+ Usage: python batch_teacher_labels.py [--workers 3]
5
+ """
6
+ import json, subprocess, sys, time, argparse
7
+ from pathlib import Path
8
+ from concurrent.futures import ThreadPoolExecutor, as_completed
9
+
10
+ ARTIFACTS = Path(__file__).resolve().parents[2] / "artifacts" / "tiny_nla"
11
+ RECORDS_FILE = ARTIFACTS / "expanded_records.jsonl"
12
+ OUTPUT = ARTIFACTS / "expanded_teacher_labels.json"
13
+ MODEL = "opencode/deepseek-v4-flash-free"
14
+
15
+ PROMPT = """\
16
+ 你是语言模型内部机制分析专家。用1-2句简洁中文描述Qwen3-0.6B在处理下面这个token时,该位置residual stream激活值编码的语义信息。
17
+
18
+ 上下文:{context}
19
+ 当前token:「{token}」(位置{pos}/{seq_len})
20
+ 模型预测下一个词的候选:{top_tokens}
21
+
22
+ 要求:简洁具体,联系token在句中的实际句法/语义角色,联系预测候选推断编码内容,不超过55字。只输出解释本身。"""
23
+
24
+
25
+ def call_opencode(record: dict) -> str:
26
+ prompt = PROMPT.format(
27
+ context=record["text"],
28
+ token=record["token_text"],
29
+ pos=record["pos"],
30
+ seq_len=record["seq_len"],
31
+ top_tokens="、".join(record["top_tokens"][:5]),
32
+ )
33
+ r = subprocess.run(
34
+ ["opencode", "run", "--model", MODEL, prompt],
35
+ capture_output=True, text=True, timeout=90,
36
+ )
37
+ lines = r.stdout.splitlines()
38
+ content = [l.strip() for l in lines
39
+ if l.strip()
40
+ and not l.strip().startswith("\x1b")
41
+ and "orchestrator" not in l
42
+ and not l.strip().startswith("{")
43
+ and not l.strip().startswith('"message"')]
44
+ return "\n".join(content).strip()
45
+
46
+
47
+ def process_record(idx_rec):
48
+ idx, rec = idx_rec
49
+ try:
50
+ expl = call_opencode(rec)
51
+ # Strip think blocks
52
+ if "<think>" in expl and "</think>" in expl:
53
+ expl = expl.split("</think>")[-1].strip()
54
+ return idx, {**rec, "teacher_explanation": expl, "teacher_source": "opencode/deepseek-v4-flash-free"}
55
+ except Exception as e:
56
+ return idx, None
57
+
58
+
59
+ def main():
60
+ parser = argparse.ArgumentParser()
61
+ parser.add_argument("--workers", type=int, default=3)
62
+ parser.add_argument("--limit", type=int, default=None)
63
+ args = parser.parse_args()
64
+
65
+ with open(RECORDS_FILE) as f:
66
+ records = [json.loads(l) for l in f]
67
+
68
+ if args.limit:
69
+ records = records[:args.limit]
70
+
71
+ # Load existing
72
+ existing = {}
73
+ if OUTPUT.exists():
74
+ with open(OUTPUT) as f:
75
+ for item in json.load(f):
76
+ existing[(item["text_idx"], item["pos"])] = item
77
+ print(f"Records: {len(records)}, existing labels: {len(existing)}")
78
+
79
+ todo = [(i, r) for i, r in enumerate(records)
80
+ if (r["text_idx"], r["pos"]) not in existing]
81
+ print(f"To process: {len(todo)}")
82
+
83
+ results = dict(existing)
84
+ done = 0
85
+
86
+ with ThreadPoolExecutor(max_workers=args.workers) as pool:
87
+ futures = {pool.submit(process_record, item): item for item in todo}
88
+ for fut in as_completed(futures):
89
+ idx, result = fut.result()
90
+ if result:
91
+ results[(result["text_idx"], result["pos"])] = result
92
+ done += 1
93
+ if done % 50 == 0 or done == len(todo):
94
+ # Checkpoint save
95
+ with open(OUTPUT, "w", encoding="utf-8") as f:
96
+ json.dump(list(results.values()), f, ensure_ascii=False, indent=2)
97
+ print(f" [{done}/{len(todo)}] saved {len(results)} labels")
98
+
99
+ with open(OUTPUT, "w", encoding="utf-8") as f:
100
+ json.dump(list(results.values()), f, ensure_ascii=False, indent=2)
101
+ print(f"\nDone: {len(results)} teacher labels → {OUTPUT}")
102
+
103
+
104
+ if __name__ == "__main__":
105
+ main()
experiments/tiny_nla/eval_final.py ADDED
@@ -0,0 +1,179 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Final model evaluation: SFT vs RL roundtrip cosine on 100 random samples."""
3
+ import os, yaml, random, argparse
4
+ from datetime import datetime
5
+
6
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
7
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
8
+ os.environ.setdefault("HF_DATASETS_OFFLINE", "1")
9
+
10
+ import torch
11
+ import torch.nn as nn
12
+ import torch.nn.functional as F
13
+ import pyarrow.parquet as pq
14
+ from pathlib import Path
15
+ from transformers import AutoModelForCausalLM, AutoTokenizer
16
+ from peft import PeftModel
17
+
18
+ # Paths
19
+ REPO_ROOT = Path(__file__).resolve().parents[2]
20
+ ARTIFACTS = REPO_ROOT / "artifacts" / "tiny_nla"
21
+ META = yaml.safe_load(open(Path(__file__).resolve().parent / "nla_meta.yaml"))
22
+ D_MODEL = META["d_model"]
23
+ INJ_CHAR = META["tokens"]["injection_char"]
24
+ INJ_TOK_ID = META["tokens"]["injection_token_id"]
25
+ INJ_SCALE = META["extraction"]["injection_scale"]
26
+ BASE_MODEL = META["base_model"]
27
+ INST_MODEL = META["av_init_model"]
28
+
29
+
30
+ class ARHead(nn.Module):
31
+ def __init__(self, d): super().__init__(); self.proj = nn.Linear(d, d, bias=False)
32
+ def forward(self, h): return F.normalize(self.proj(h), dim=-1)
33
+
34
+
35
+ def tprint(s): print(f"[{datetime.now().strftime('%H:%M:%S')}] {s}", flush=True)
36
+
37
+
38
+ def load_activations():
39
+ t = pq.read_table(ARTIFACTS / "activations_v2.parquet")
40
+ acts = torch.tensor([t["activation"][i].as_py() for i in range(len(t))], dtype=torch.float32)
41
+ tprint(f"Loaded {len(acts)} activations")
42
+ return acts
43
+
44
+
45
+ def load_av_model(path, device):
46
+ tprint(f"Loading AV from {path}...")
47
+ tok = AutoTokenizer.from_pretrained(INST_MODEL)
48
+ base = AutoModelForCausalLM.from_pretrained(
49
+ INST_MODEL, trust_remote_code=True, dtype=torch.float16,
50
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(device)
51
+ model = PeftModel.from_pretrained(base, path)
52
+ model.eval()
53
+ return model, tok
54
+
55
+
56
+ def load_ar(device):
57
+ ckpt = torch.load(ARTIFACTS / "checkpoints" / "ar_v2" / "ar_head_v2.pt",
58
+ map_location=device, weights_only=True)
59
+ head = ARHead(D_MODEL).to(device)
60
+ head.load_state_dict(ckpt["head"])
61
+ head.eval()
62
+ tprint(f"AR head loaded (val_cos={ckpt.get('val_cosine',0):.4f})")
63
+ return head
64
+
65
+
66
+ def load_ar_backbone(device):
67
+ tprint(f"Loading AR backbone ({BASE_MODEL})...")
68
+ m = AutoModelForCausalLM.from_pretrained(
69
+ BASE_MODEL, trust_remote_code=True, dtype=torch.float16,
70
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(device)
71
+ m.eval()
72
+ return m
73
+
74
+
75
+ def generate(model, tok, act_scaled, max_new=64):
76
+ prompt = f"<concept>{INJ_CHAR}</concept>\n<explanation>"
77
+ p_ids = tok(prompt, return_tensors="pt")["input_ids"].to(act_scaled.device)
78
+ p_mask = torch.ones(1, p_ids.shape[1], device=act_scaled.device, dtype=torch.long)
79
+ inj_pos = (p_ids[0] == INJ_TOK_ID).nonzero(as_tuple=True)[0][0].item()
80
+ embeds = model.get_input_embeddings()(p_ids).clone()
81
+ embeds[0, inj_pos] = act_scaled[0].to(embeds.dtype)
82
+ with torch.no_grad():
83
+ out = model.generate(inputs_embeds=embeds, attention_mask=p_mask,
84
+ max_new_tokens=max_new, do_sample=False,
85
+ pad_token_id=tok.eos_token_id)
86
+ return tok.decode(out[0], skip_special_tokens=True).strip()
87
+
88
+
89
+ def reconstruct(ar_backbone, ar_head, tok, explanation):
90
+ tok.pad_token_id = tok.eos_token_id
91
+ dev = ar_head.proj.weight.device
92
+ enc = tok([explanation], return_tensors="pt", padding=True,
93
+ truncation=True, max_length=128).to(dev)
94
+ with torch.no_grad():
95
+ h = ar_backbone(**enc, output_hidden_states=True).hidden_states[-1]
96
+ lens = enc["attention_mask"].sum(1) - 1
97
+ last = h[0, lens[0]]
98
+ recon = ar_head(last.unsqueeze(0))
99
+ return recon.float()
100
+
101
+
102
+ def eval_one(model, tok, act_raw, ar_backbone, ar_head, tok_ar):
103
+ act_s = act_raw.unsqueeze(0) / act_raw.norm() * INJ_SCALE
104
+ act_s = act_s.to(ar_head.proj.weight.device)
105
+ act_n = (act_raw.unsqueeze(0) / act_raw.norm()).to(ar_head.proj.weight.device)
106
+ expl = generate(model, tok, act_s)
107
+ recon = reconstruct(ar_backbone, ar_head, tok_ar, expl)
108
+ return (recon * act_n).sum(-1).item(), expl
109
+
110
+
111
+ # ══════════════════════════════════════════════════════════
112
+ if __name__ == "__main__":
113
+ parser = argparse.ArgumentParser()
114
+ parser.add_argument("--num-samples", type=int, default=100)
115
+ parser.add_argument("--seed", type=int, default=42)
116
+ args = parser.parse_args()
117
+
118
+ device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
119
+ random.seed(args.seed)
120
+
121
+ print(f"\n{'='*55}")
122
+ print(f" Tiny-NLA Final Eval | samples={args.num_samples} | seed={args.seed}")
123
+ print(f"{'='*55}\n")
124
+
125
+ acts = load_activations()
126
+ ar_head = load_ar(device)
127
+ ar_bb = load_ar_backbone(device)
128
+ tok_ar = AutoTokenizer.from_pretrained(BASE_MODEL)
129
+
130
+ sft_model, sft_tok = load_av_model(ARTIFACTS / "checkpoints" / "av_v2", device)
131
+ rl_model, rl_tok = load_av_model(ARTIFACTS / "checkpoints" / "av_rl_best", device)
132
+
133
+ eval_idx = random.sample(range(len(acts)), args.num_samples)
134
+
135
+ sft_cos, rl_cos = [], []
136
+ sft_expls, rl_expls = [], []
137
+
138
+ tprint(f"Evaluating {args.num_samples} samples...")
139
+ for k, idx in enumerate(eval_idx):
140
+ act = acts[idx]
141
+ c1, e1 = eval_one(sft_model, sft_tok, act, ar_bb, ar_head, tok_ar)
142
+ c2, e2 = eval_one(rl_model, rl_tok, act, ar_bb, ar_head, tok_ar)
143
+ sft_cos.append(c1); rl_cos.append(c2)
144
+ sft_expls.append(e1); rl_expls.append(e2)
145
+ if (k + 1) % 20 == 0:
146
+ tprint(f" {k+1}/{args.num_samples} | sft={sum(sft_cos)/(k+1):.4f} | rl={sum(rl_cos)/(k+1):.4f}")
147
+
148
+ sft_mean = sum(sft_cos) / len(sft_cos)
149
+ rl_mean = sum(rl_cos) / len(rl_cos)
150
+ delta = rl_mean - sft_mean
151
+ gains = [rl_cos[i] - sft_cos[i] for i in range(args.num_samples)]
152
+ positive = sum(1 for g in gains if g > 0)
153
+
154
+ print(f"\n{'='*55}")
155
+ print(f" RESULTS")
156
+ print(f"{'='*55}")
157
+ print(f" {'':16} {'SFT':>10} {'RL':>10} {'Δ':>10}")
158
+ print(f" {'Mean':16} {sft_mean:10.4f} {rl_mean:10.4f} {delta:+10.4f}")
159
+ print(f" {'Best':16} {max(sft_cos):10.4f} {max(rl_cos):10.4f} {max(rl_cos)-max(sft_cos):+10.4f}")
160
+ print(f" {'Worst':16} {min(sft_cos):10.4f} {min(rl_cos):10.4f} {min(rl_cos)-min(sft_cos):+10.4f}")
161
+ print(f" {'RL wins':16} {positive}/{args.num_samples} ({100*positive/args.num_samples:.0f}%)")
162
+ print(f"{'='*55}")
163
+
164
+ # Top 5 improvements
165
+ print(f"\n── TOP 5 GAINS (RL − SFT) ──")
166
+ sorted_idx = sorted(range(args.num_samples), key=lambda i: gains[i], reverse=True)
167
+ for rank, i in enumerate(sorted_idx[:5]):
168
+ print(f"\n #{rank+1} Δ={gains[i]:+.4f} | SFT cos={sft_cos[i]:.4f}")
169
+ print(f" SFT: {sft_expls[i][:130]}")
170
+ print(f" RL: {rl_expls[i][:130]}")
171
+
172
+ # Bottom 5
173
+ print(f"\n── BOTTOM 5 (RL regressions) ──")
174
+ for rank, i in enumerate(sorted_idx[-5:]):
175
+ print(f"\n #{args.num_samples-4+rank} Δ={gains[i]:+.4f} | RL cos={rl_cos[i]:.4f}")
176
+ print(f" SFT: {sft_expls[i][:130]}")
177
+ print(f" RL: {rl_expls[i][:130]}")
178
+
179
+ print(f"\n{tprint('Done.')}")
experiments/tiny_nla/eval_roundtrip.py ADDED
@@ -0,0 +1,157 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Round-trip evaluation: extract → explain → reconstruct → metrics
4
+ Runs on held-out data and produces a comprehensive report.
5
+ """
6
+
7
+ import json, sys, yaml, random
8
+ from pathlib import Path
9
+
10
+ import torch
11
+ import torch.nn.functional as F
12
+
13
+ sys.path.insert(0, str(Path(__file__).resolve().parent))
14
+ from infer_tiny_nla import TinyNLA
15
+
16
+
17
+ REPO_ROOT = Path(__file__).resolve().parents[2]
18
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
19
+
20
+
21
+ def main():
22
+ print("=" * 60)
23
+ print("🔄 Tiny-NLA Round-Trip Evaluation")
24
+ print("=" * 60)
25
+
26
+ nla = TinyNLA()
27
+
28
+ # Load dataset for held-out samples
29
+ with open(ARTIFACTS_DIR / "dataset.jsonl", "r", encoding="utf-8") as f:
30
+ records = [json.loads(l) for l in f]
31
+
32
+ activations = torch.load(ARTIFACTS_DIR / "activations.pt", weights_only=True)
33
+
34
+ # Filter to valid explanations
35
+ valid = []
36
+ for i, r in enumerate(records):
37
+ exp = r.get("teacher_explanation", "") or r.get("teacher_explanation_raw", "")
38
+ if exp and exp not in ("[空输出]", ""):
39
+ valid.append((r, activations[i]))
40
+
41
+ print(f" Total records: {len(records)}, Valid: {len(valid)}")
42
+
43
+ # Use a set of held-out indices for evaluation
44
+ # We'll use the first N samples from different texts
45
+ # Stratify by text to ensure diversity
46
+ texts_grouped = {}
47
+ for i, (r, _) in enumerate(valid):
48
+ texts_grouped.setdefault(r["text_idx"], []).append(i)
49
+
50
+ held_out = []
51
+ for tidx, indices in texts_grouped.items():
52
+ # Take last 2 from each text group as held-out
53
+ held_out.extend(indices[-2:])
54
+
55
+ # Make sure we have at least 20 held-out
56
+ if len(held_out) < 20:
57
+ extra = [i for i in range(len(valid)) if i not in held_out]
58
+ random.Random(42).shuffle(extra)
59
+ held_out.extend(extra[:20 - len(held_out)])
60
+
61
+ print(f" Held-out samples: {len(held_out)}")
62
+
63
+ # Round-trip evaluation
64
+ results = []
65
+ for idx in held_out:
66
+ rec, act = valid[idx]
67
+
68
+ # We need original text + position for context
69
+ text = rec["text"]
70
+ pos = rec["pos"]
71
+ token_text = rec["token_text"]
72
+ teacher_exp = rec.get("teacher_explanation", "") or rec.get("teacher_explanation_raw", "")
73
+
74
+ # Run round-trip
75
+ try:
76
+ # Extract activation
77
+ ext = nla.extract(text, pos)
78
+
79
+ # Generate AV explanation (from the activation, not the text)
80
+ av_result = nla.explain(ext["activation"])
81
+ av_explanation = av_result["explanation"]
82
+
83
+ # Reconstruct from AV output
84
+ rec_result = nla.reconstruct(av_explanation)
85
+
86
+ row = {
87
+ "text": text,
88
+ "position": pos,
89
+ "token_text": token_text,
90
+ "teacher_explanation": teacher_exp[:120],
91
+ "av_explanation": av_explanation[:120],
92
+ }
93
+
94
+ if "reconstructed" in rec_result:
95
+ orig_n = F.normalize(act.unsqueeze(0), dim=-1)
96
+ recon = rec_result["reconstructed"].unsqueeze(0)
97
+ recon_n = F.normalize(recon, dim=-1)
98
+ cosine = (orig_n * recon_n).sum(dim=-1).item()
99
+ mse = F.mse_loss(orig_n, recon_n).item()
100
+ row["roundtrip_cosine"] = round(cosine, 4)
101
+ row["roundtrip_mse"] = round(mse, 6)
102
+
103
+ # Also compute: teacher_explanation → original_activation cosine
104
+ tea_result = nla.reconstruct(teacher_exp)
105
+ if "reconstructed" in tea_result:
106
+ tea_recon = tea_result["reconstructed"].unsqueeze(0)
107
+ tea_recon_n = F.normalize(tea_recon, dim=-1)
108
+ tea_cosine = (orig_n * tea_recon_n).sum(dim=-1).item()
109
+ tea_mse = F.mse_loss(orig_n, tea_recon_n).item()
110
+ row["teacher_to_activation_cosine"] = round(tea_cosine, 4)
111
+ row["teacher_to_activation_mse"] = round(tea_mse, 6)
112
+
113
+ results.append(row)
114
+
115
+ except Exception as e:
116
+ print(f" ⚠️ Error at idx {idx}: {e}")
117
+ continue
118
+
119
+ # Summary statistics
120
+ rt_cosines = [r.get("roundtrip_cosine", 0) for r in results if "roundtrip_cosine" in r]
121
+ tea_cosines = [r.get("teacher_to_activation_cosine", 0) for r in results if "teacher_to_activation_cosine" in r]
122
+
123
+ print(f"\n📊 Round-Trip Metrics")
124
+ print(f" Samples evaluated: {len(results)}")
125
+ if rt_cosines:
126
+ print(f" Round-trip (AV→AR) cosine:")
127
+ print(f" Mean: {sum(rt_cosines)/len(rt_cosines):.4f}")
128
+ print(f" Min: {min(rt_cosines):.4f}")
129
+ print(f" Max: {max(rt_cosines):.4f}")
130
+ if tea_cosines:
131
+ print(f" Teacher→Activation cosine (AR upper bound):")
132
+ print(f" Mean: {sum(tea_cosines)/len(tea_cosines):.4f}")
133
+ print(f" Min: {min(tea_cosines):.4f}")
134
+ print(f" Max: {max(tea_cosines):.4f}")
135
+
136
+ # Save results
137
+ out_path = ARTIFACTS_DIR / "roundtrip_results.json"
138
+ with open(out_path, "w", encoding="utf-8") as f:
139
+ json.dump(results, f, ensure_ascii=False, indent=2)
140
+ print(f"\n Results saved: {out_path}")
141
+
142
+ # Show 20 worked examples
143
+ print(f"\n📝 20 Worked Examples")
144
+ print("=" * 60)
145
+ for i, r in enumerate(results[:20]):
146
+ print(f"\n [{i+1}] Text: {r['text'][:50]}...")
147
+ print(f" Token: {r['token_text']!r} (pos={r['position']})")
148
+ print(f" Teacher: {r['teacher_explanation'][:80]}")
149
+ print(f" AV: {r['av_explanation'][:80]}")
150
+ if "roundtrip_cosine" in r:
151
+ print(f" Round-trip cos: {r['roundtrip_cosine']}")
152
+ if "teacher_to_activation_cosine" in r:
153
+ print(f" Teacher→Act cos: {r['teacher_to_activation_cosine']}")
154
+
155
+
156
+ if __name__ == "__main__":
157
+ main()
experiments/tiny_nla/expand_dataset.py ADDED
@@ -0,0 +1,82 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Generate diverse Chinese texts via opencode, then extract activations and teacher labels.
4
+ Run: python expand_dataset.py
5
+ """
6
+ import json, subprocess, sys, time
7
+ from pathlib import Path
8
+
9
+ REPO_ROOT = Path(__file__).resolve().parents[2]
10
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
11
+ OUT_TEXTS = ARTIFACTS_DIR / "expanded_texts.json"
12
+ MODEL = "opencode/deepseek-v4-flash-free"
13
+
14
+ DOMAINS = [
15
+ ("科技AI", 20), ("日常生活", 20), ("自然科学", 20), ("社会新闻", 20),
16
+ ("文学文化", 20), ("经济商业", 20), ("历史哲学", 20), ("医学健康", 15),
17
+ ("体育运动", 15), ("环境生态", 15), ("教育学习", 15), ("人际关系", 15),
18
+ ("法律政治", 10), ("饮食美食", 10), ("旅游地理", 10),
19
+ ]
20
+
21
+ PROMPT_TEMPLATE = """\
22
+ 请生成 {n} 条不同的中文句子,主题是「{domain}」。
23
+
24
+ 要求:
25
+ - 每条 15-50 个字,完整句子
26
+ - 内容多样,涵盖该主题的不同角度
27
+ - 语言自然,包含实词(名词、动词、形容词)和虚词(介词、连词、语气词)
28
+ - 句式多样:陈述句、疑问句、复合句都可以
29
+ - 不要编号,每行一句,直接输出句子
30
+
31
+ 只输出句子,不要其他解释。"""
32
+
33
+
34
+ def call_opencode(prompt: str) -> str:
35
+ r = subprocess.run(
36
+ ["opencode", "run", "--model", MODEL, prompt],
37
+ capture_output=True, text=True, timeout=120,
38
+ )
39
+ lines = r.stdout.splitlines()
40
+ content = [l for l in lines if l.strip()
41
+ and not l.strip().startswith("\x1b")
42
+ and "> orchestrator" not in l
43
+ and not l.strip().startswith("{")
44
+ and not l.strip().startswith('"')]
45
+ return "\n".join(content).strip()
46
+
47
+
48
+ def main():
49
+ existing = []
50
+ if OUT_TEXTS.exists():
51
+ with open(OUT_TEXTS) as f:
52
+ existing = json.load(f)
53
+ print(f"Existing texts: {len(existing)}")
54
+
55
+ existing_set = set(existing)
56
+ all_texts = list(existing)
57
+
58
+ for domain, n in DOMAINS:
59
+ prompt = PROMPT_TEMPLATE.format(domain=domain, n=n)
60
+ print(f"\n [{domain}] requesting {n} sentences...")
61
+ try:
62
+ out = call_opencode(prompt)
63
+ sentences = [l.strip() for l in out.splitlines()
64
+ if l.strip() and 8 <= len(l.strip()) <= 80
65
+ and l.strip() not in existing_set]
66
+ print(f" Got {len(sentences)} new sentences")
67
+ for s in sentences[:5]:
68
+ print(f" {s[:50]}")
69
+ all_texts.extend(sentences)
70
+ existing_set.update(sentences)
71
+ time.sleep(0.5)
72
+ except Exception as e:
73
+ print(f" Error: {e}", file=sys.stderr)
74
+
75
+ # Save
76
+ with open(OUT_TEXTS, "w", encoding="utf-8") as f:
77
+ json.dump(all_texts, f, ensure_ascii=False, indent=2)
78
+ print(f"\nTotal texts: {len(all_texts)} → {OUT_TEXTS}")
79
+
80
+
81
+ if __name__ == "__main__":
82
+ main()
experiments/tiny_nla/expand_dataset_v2.py ADDED
@@ -0,0 +1,61 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Phase 0.1: expand Chinese text corpus to 500+ sentences via opencode (parallel)."""
3
+ import json, subprocess, sys, time
4
+ from pathlib import Path
5
+ from concurrent.futures import ThreadPoolExecutor, as_completed
6
+
7
+ ARTIFACTS = Path(__file__).resolve().parents[2] / "artifacts" / "tiny_nla"
8
+ OUT = ARTIFACTS / "expanded_texts.json"
9
+ MODEL = "opencode/deepseek-v4-flash-free"
10
+
11
+ # Additional domains to push past 500
12
+ DOMAINS = [
13
+ ("编程技术", 25), ("数学逻辑", 20), ("影视娱乐", 20), ("宗教信仰", 15),
14
+ ("语言学", 20), ("艺术美学", 20), ("交通出行", 20), ("宠物动物", 20),
15
+ ("天气气候", 15), ("金融投资", 20), ("职场沟通", 20), ("心理学", 20),
16
+ ("社交媒体", 20), ("游戏娱乐", 15), ("亲子教育", 15),
17
+ ]
18
+
19
+
20
+ def gen_sentences(domain_n):
21
+ domain, n = domain_n
22
+ # Request a mix of short/medium/long
23
+ prompt = (
24
+ f"请生成{n}条关于「{domain}」的中文句子。"
25
+ "要求:包含8-15字短句、15-30字中等句、30-50字长句各占约三分之一;"
26
+ "句式多样(陈述、疑问、复合);语言自然,含丰富实词和虚词。"
27
+ "不要编号,每行一句,只输出句子。"
28
+ )
29
+ r = subprocess.run(["opencode", "run", "--model", MODEL, prompt],
30
+ capture_output=True, text=True, timeout=120)
31
+ lines = r.stdout.splitlines()
32
+ sents = [l.strip() for l in lines
33
+ if l.strip() and 5 <= len(l.strip()) <= 80
34
+ and not l.strip().startswith("\x1b")
35
+ and "orchestrator" not in l
36
+ and not l.strip().startswith("{")]
37
+ return domain, sents
38
+
39
+
40
+ def main():
41
+ with open(OUT) as f:
42
+ existing = json.load(f)
43
+ existing_set = set(existing)
44
+ print(f"Starting with {len(existing)} sentences")
45
+
46
+ with ThreadPoolExecutor(max_workers=3) as pool:
47
+ futures = {pool.submit(gen_sentences, d): d for d in DOMAINS}
48
+ for fut in as_completed(futures):
49
+ domain, sents = fut.result()
50
+ new = [s for s in sents if s not in existing_set]
51
+ existing.extend(new)
52
+ existing_set.update(new)
53
+ print(f" [{domain}] +{len(new)} → total {len(existing)}")
54
+
55
+ with open(OUT, "w", encoding="utf-8") as f:
56
+ json.dump(existing, f, ensure_ascii=False, indent=2)
57
+ print(f"\nFinal corpus: {len(existing)} sentences → {OUT}")
58
+
59
+
60
+ if __name__ == "__main__":
61
+ main()
experiments/tiny_nla/gen_demo_html.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """Generate a standalone HTML demo from real teacher labels."""
3
+ import json
4
+ from pathlib import Path
5
+ from collections import defaultdict
6
+
7
+ LABELS_FILE = Path("/Users/cccmmd/InfoLens/artifacts/tiny_nla/teacher_labels_hq.json")
8
+ OUTPUT_HTML = Path("/Users/cccmmd/InfoLens/experiments/tiny_nla/demo_activation_translator.html")
9
+
10
+ with open(LABELS_FILE) as f:
11
+ labels = json.load(f)
12
+
13
+ by_text = defaultdict(list)
14
+ for r in labels:
15
+ by_text[r["text_idx"]].append(r)
16
+
17
+ sentences = []
18
+ for tid in sorted(by_text):
19
+ recs = sorted(by_text[tid], key=lambda x: x["pos"])
20
+ sent = recs[0]["text"]
21
+ tokens = []
22
+ for r in recs:
23
+ tokens.append({
24
+ "text": r["token_text"],
25
+ "pos": r["pos"],
26
+ "top5": r["top_tokens"],
27
+ "explanation": r["teacher_explanation"],
28
+ "norm": round(r["activation_norm"], 1),
29
+ })
30
+ sentences.append({"text": sent, "tokens": tokens})
31
+ if len(sentences) >= 5:
32
+ break
33
+
34
+ sentences_json = json.dumps(sentences, ensure_ascii=False)
35
+
36
+ # Read template
37
+ import os
38
+ template_path = Path(__file__).parent / "demo_template.html"
39
+ html = template_path.read_text(encoding="utf-8")
40
+ html = html.replace("__SENTENCES_DATA__", sentences_json)
41
+ OUTPUT_HTML.write_text(html, encoding="utf-8")
42
+ print(f"Demo written to {OUTPUT_HTML}")
43
+ print(f"Sentences: {len(sentences)}, tokens: {sum(len(s['tokens']) for s in sentences)}")
experiments/tiny_nla/gen_opencode_only.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """opencode-only teacher label generation - runs in background."""
3
+ import json, subprocess, time, sys
4
+ from pathlib import Path
5
+ from concurrent.futures import ThreadPoolExecutor, as_completed
6
+
7
+ ARTIFACTS = Path(__file__).resolve().parents[2] / "artifacts" / "tiny_nla"
8
+ RECORDS_FILE = ARTIFACTS / "records_v2.jsonl"
9
+ OUTPUT_JSONL = ARTIFACTS / "teacher_labels_v2.jsonl"
10
+
11
+ PROMPT = """\
12
+ 你是语言模型内部机制分析专家。用1-2句简洁中文描述Qwen3-0.6B在处理下面这个token时,该位置residual stream激活值编码的语义信息。
13
+
14
+ 上下文:{context}
15
+ 当前token:「{token}」(位置{pos}/{seq_len})
16
+ 模型预测下一个词的候选:{top_tokens}
17
+
18
+ 要求:简洁具体,联系token在句中的实际句法/语义角色,联系预测候选推断编码内容,不超过55字。只输出解释本身。"""
19
+
20
+
21
+ def make_prompt(r):
22
+ return PROMPT.format(context=r["text"], token=r["token_text"],
23
+ pos=r["pos"], seq_len=r["seq_len"],
24
+ top_tokens="、".join(r["top_tokens"][:5]))
25
+
26
+
27
+ def strip_think(t):
28
+ return t.split("</think>")[-1].strip() if "</think>" in t else t.strip()
29
+
30
+
31
+ def call_opencode(r):
32
+ res = subprocess.run(
33
+ ["opencode", "run", "--model", "opencode/deepseek-v4-flash-free", make_prompt(r)],
34
+ capture_output=True, text=True, timeout=240,
35
+ )
36
+ lines = [l.strip() for l in res.stdout.splitlines()
37
+ if l.strip() and not l.startswith("\x1b") and "orchestrator" not in l
38
+ and not l.startswith("{") and not l.startswith('"')]
39
+ return strip_think("\n".join(lines))
40
+
41
+
42
+ def load_done():
43
+ done = set()
44
+ try:
45
+ for r in json.load(open(ARTIFACTS / "teacher_labels_v2.json")):
46
+ done.add((r["text_idx"], r["pos"]))
47
+ except Exception:
48
+ pass
49
+ if OUTPUT_JSONL.exists():
50
+ for line in open(OUTPUT_JSONL):
51
+ try:
52
+ r = json.loads(line)
53
+ done.add((r["text_idx"], r["pos"]))
54
+ except Exception:
55
+ pass
56
+ return done
57
+
58
+
59
+ def process(idx_rec):
60
+ idx, rec = idx_rec
61
+ try:
62
+ expl = call_opencode(rec)
63
+ if expl and len(expl) >= 5:
64
+ return {**rec, "teacher_explanation": expl, "teacher_source": "opencode"}
65
+ except Exception as e:
66
+ print(f" err [{idx}]: {e}", flush=True)
67
+ return None
68
+
69
+
70
+ def main():
71
+ workers = int(sys.argv[1]) if len(sys.argv) > 1 else 3
72
+ records = [json.loads(l) for l in open(RECORDS_FILE)]
73
+ done = load_done()
74
+ todo = [(i, r) for i, r in enumerate(records)
75
+ if (r["text_idx"], r["pos"]) not in done]
76
+ print(f"Remaining: {len(todo)}/9810, workers: {workers}", flush=True)
77
+
78
+ completed = 0
79
+ t0 = time.time()
80
+ out_f = open(OUTPUT_JSONL, "a", encoding="utf-8", buffering=1)
81
+
82
+ with ThreadPoolExecutor(max_workers=workers) as pool:
83
+ futures = {pool.submit(process, item): item for item in todo}
84
+ for fut in as_completed(futures):
85
+ result = fut.result()
86
+ if result:
87
+ out_f.write(json.dumps(result, ensure_ascii=False) + "\n")
88
+ completed += 1
89
+ if completed % 20 == 0:
90
+ rate = completed / (time.time() - t0) * 60
91
+ eta_h = (len(todo) - completed) / (rate / 60) / 3600
92
+ print(f" [{len(done)+completed}/{9810}] {rate:.0f}/min ETA={eta_h:.1f}h", flush=True)
93
+
94
+ out_f.close()
95
+ print(f"Done: {completed} new labels written", flush=True)
96
+
97
+
98
+ if __name__ == "__main__":
99
+ main()
experiments/tiny_nla/generate_data.py ADDED
@@ -0,0 +1,397 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Phase 3: Data Generation
4
+
5
+ Two-pass process:
6
+ 1. Load BASE model → extract activations for token positions in diverse texts
7
+ 2. Load INSTRUCT model → generate teacher explanations
8
+
9
+ Output: artifacts/tiny_nla/dataset.jsonl + artifacts/tiny_nla/dataset_stats.json
10
+ """
11
+
12
+ import os, sys, json, math, yaml, time, random
13
+ from pathlib import Path
14
+
15
+ import torch
16
+ from transformers import AutoModelForCausalLM, AutoTokenizer
17
+
18
+ # ── paths ──────────────────────────────────────────────
19
+ REPO_ROOT = Path(__file__).resolve().parents[2]
20
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
21
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
22
+ ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
23
+
24
+ BASE_MODEL = "Qwen/Qwen3-0.6B-Base"
25
+ INSTRUCT_MODEL = "Qwen/Qwen3-0.6B"
26
+
27
+ # ── device ─────────────────────────────────────────────
28
+ def detect_device():
29
+ if torch.cuda.is_available():
30
+ return torch.device("cuda")
31
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
32
+ return torch.device("mps")
33
+ return torch.device("cpu")
34
+
35
+ def get_dtype(device):
36
+ if device.type == "cuda":
37
+ return torch.float16
38
+ return torch.float32 # MPS/CPU safest
39
+
40
+
41
+ # ── diverse Chinese texts ──────────────────────────────
42
+ TEXTS = [
43
+ # Technology & AI
44
+ "人工智能正在深刻改变各行各业的发展模式,从自动驾驶到医疗诊断,都在不断突破。",
45
+ "深度学习模型通过多层神经网络学习数据的层次化特征表示,在自然语言处理领域取得了显著成果。",
46
+ "量子计算是一种利用量子力学原理处理信息的新型计算范式。",
47
+
48
+ # Daily life
49
+ "今天天气真好,阳光明媚,我们一起去公园散步吧,感受大自然的美丽。",
50
+ "这家餐厅的菜非常好吃,特别是他们的招牌菜红烧肉,味道鲜美极了。",
51
+ "昨天晚上我看了一部很感人的电影,讲述了一个关于友情和梦想的故事。",
52
+
53
+ # Science & education
54
+ "在学习编程的过程中,最重要的是掌握基本的逻辑思维和解决问题的能力。",
55
+ "数学是自然科学的基础,它帮助我们理解世界的规律和结构。",
56
+ "历史上许多伟大的科学家都经历了无数次失败才取得了突破性的发现。",
57
+
58
+ # News & society
59
+ "近年来,全球气候变化问题日益受到各国政府的高度关注。",
60
+ "数字经济正在成为推动全球经济增长的新引擎。",
61
+ "教育公平是社会公平的重要基础,需要全社会共同努力。",
62
+
63
+ # Literature & culture
64
+ "读书是一种与作者对话的方式,通过阅读我们可以获得知识和智慧。",
65
+ "中国传统文化的魅力在于其深厚的历史积淀和独特的哲学思想。",
66
+ "音乐是人类共同的语言,它能够跨越国界传达情感。",
67
+
68
+ # Economy & business
69
+ "市场经济的核心在于供需关系的动态平衡。",
70
+ "创新是企业保持竞争力的关键因素。",
71
+ "投资理财需要长期规划和理性决策。",
72
+
73
+ # Short phrases
74
+ "请帮我解释一下这个概念。",
75
+ "我想了解更多关于这个话题的信息。",
76
+ "这是一个非常重要的发现。",
77
+ "我们需要认真对待这个问题。",
78
+ ]
79
+
80
+
81
+ # ── Pass 1: Extract activations ─────────────────────────
82
+ def extract_activations(model, tokenizer, texts, layer_idx, device, dtype):
83
+ """
84
+ For each text, extract residual stream activation at layer_idx
85
+ for EVERY token position (not just last). Returns list of records.
86
+ """
87
+ print("=" * 60)
88
+ print("📥 Pass 1: Extracting Activations")
89
+ print("=" * 60)
90
+ records = []
91
+ total_tokens = 0
92
+
93
+ for idx, text in enumerate(texts):
94
+ inputs = tokenizer(text, return_tensors="pt").to(device)
95
+ input_ids = inputs["input_ids"] # [1, seq_len]
96
+ seq_len = input_ids.shape[1]
97
+
98
+ with torch.no_grad():
99
+ outputs = model(
100
+ **inputs,
101
+ output_hidden_states=True,
102
+ )
103
+
104
+ # Get logits for last token position for top-k info
105
+ logits = outputs.logits[0] # [seq_len, vocab]
106
+
107
+ # Get hidden states for our target layer
108
+ hidden = outputs.hidden_states[layer_idx] # [1, seq_len, d_model]
109
+
110
+ for pos in range(seq_len):
111
+ token_id = input_ids[0, pos].item()
112
+ token_text = tokenizer.decode([token_id])
113
+ activation = hidden[0, pos, :].cpu() # [d_model]
114
+
115
+ # Top-k at this position (from logits at this position predicting NEXT token)
116
+ if pos < seq_len - 1:
117
+ next_logits = logits[pos] # logits for predicting next token after pos
118
+ topk_vals, topk_idxs = torch.topk(next_logits, k=10)
119
+ top_tokens = [tokenizer.decode([t]) for t in topk_idxs]
120
+ else:
121
+ # Last position has no "next token" prediction in base model
122
+ # But we still extract the activation for it
123
+ next_logits = logits[pos]
124
+ topk_vals, topk_idxs = torch.topk(next_logits, k=10)
125
+ top_tokens = [tokenizer.decode([t]) for t in topk_idxs]
126
+
127
+ records.append({
128
+ "text_idx": idx,
129
+ "text": text,
130
+ "pos": pos,
131
+ "token_id": token_id,
132
+ "token_text": token_text,
133
+ "activation_vector": activation.tolist(),
134
+ "top_tokens": top_tokens,
135
+ "activation_norm": activation.norm().item(),
136
+ })
137
+ total_tokens += 1
138
+
139
+ if (idx + 1) % 5 == 0:
140
+ print(f" Processed {idx + 1}/{len(texts)} texts ({total_tokens} tokens so far)")
141
+
142
+ print(f" ✅ Extracted {total_tokens} activations from {len(texts)} texts")
143
+ return records
144
+
145
+
146
+ # ── Pass 2: Generate teacher explanations ──────────────
147
+ def generate_teacher_explanations(
148
+ model, tokenizer, records, device, dtype,
149
+ max_samples=500, batch_prompt=False,
150
+ ):
151
+ """
152
+ For each record, generate a teacher explanation using the instruct model.
153
+
154
+ Teacher prompt (instruct format):
155
+ Given a context and token position, explain what the model might be focusing on.
156
+ """
157
+ print("=" * 60)
158
+ print("📝 Pass 2: Generating Teacher Explanations")
159
+ print("=" * 60)
160
+
161
+ # Determine instruct chat template format
162
+ # Qwen3 instruct uses: <|im_start|>system...<|im_end|> etc
163
+ # But for base model style, just use a simple prompt
164
+
165
+ system_msg = "你是一个模型可解释性专家。给定一段文本和其中的某个位置,用1-2句中文简短解释模型在该位置可能关注什么语义信息。不要长篇分析。"
166
+
167
+ # Sample records, ensuring diversity across texts
168
+ random.seed(42)
169
+ n_texts = len(set(r["text_idx"] for r in records))
170
+ samples_per_text = max(1, max_samples // n_texts)
171
+
172
+ # Stratified sample: take at least 15 tokens per text (or all if fewer available)
173
+ texts_grouped = {}
174
+ for r in records:
175
+ texts_grouped.setdefault(r["text_idx"], []).append(r)
176
+
177
+ sampled_records = []
178
+ for tidx in sorted(texts_grouped.keys()):
179
+ group = texts_grouped[tidx]
180
+ # Take evenly spaced positions for diversity
181
+ n = min(len(group), max(15, samples_per_text))
182
+ indices = sorted(random.sample(range(len(group)), min(n, len(group))))
183
+ for i in indices:
184
+ sampled_records.append(group[i])
185
+
186
+ # If we have fewer than target, take more
187
+ if len(sampled_records) < max_samples and len(sampled_records) < len(records):
188
+ existing_ids = set(id(r) for r in sampled_records)
189
+ more = [r for r in records if id(r) not in existing_ids]
190
+ remaining = max_samples - len(sampled_records)
191
+ sampled_records.extend(random.sample(more, min(remaining, len(more))))
192
+
193
+ print(f" Target: {max_samples} samples, selected {len(sampled_records)} for labeling")
194
+
195
+ results = []
196
+ teacher_failures = 0
197
+
198
+ for i, rec in enumerate(sampled_records):
199
+ text = rec["text"]
200
+ pos = rec["pos"]
201
+ token_text = rec["token_text"]
202
+ top_tokens = rec["top_tokens"]
203
+
204
+ # Build instruction prompt
205
+ user_msg = (
206
+ f"文本:{text}\n"
207
+ f"位置:第{pos}个token(文本位置),该token文本是「{token_text}」\n"
208
+ f"该位置模型预测的下一个token候选:{', '.join(top_tokens[:5])}\n\n"
209
+ f"请用1-2句中文解释:模型在这个位置可能关注什么?"
210
+ )
211
+
212
+ messages = [
213
+ {"role": "system", "content": system_msg},
214
+ {"role": "user", "content": user_msg},
215
+ ]
216
+
217
+ prompt = tokenizer.apply_chat_template(
218
+ messages, tokenize=False, add_generation_prompt=True
219
+ )
220
+
221
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
222
+
223
+ try:
224
+ with torch.no_grad():
225
+ output_ids = model.generate(
226
+ **inputs,
227
+ max_new_tokens=64,
228
+ do_sample=True,
229
+ temperature=0.7,
230
+ top_p=0.9,
231
+ pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
232
+ )
233
+ # Extract the generated part (skip input prompt)
234
+ generated = output_ids[0][inputs["input_ids"].shape[1]:]
235
+ explanation = tokenizer.decode(generated, skip_special_tokens=True).strip()
236
+
237
+ if not explanation or len(explanation) < 5:
238
+ explanation = "[空输出]"
239
+ teacher_failures += 1
240
+
241
+ except Exception as e:
242
+ explanation = f"[生成失败: {e}]"
243
+ teacher_failures += 1
244
+
245
+ rec["teacher_explanation"] = explanation
246
+ results.append(rec)
247
+
248
+ if (i + 1) % 20 == 0 or i == 0:
249
+ print(f" [{i+1}/{len(sampled_records)}] token={token_text!r} -> {explanation[:60]}...")
250
+
251
+ print(f"\n ✅ Generated {len(results)} explanations")
252
+ print(f" ⚠️ Failures/empty: {teacher_failures}/{len(results)}")
253
+ return results
254
+
255
+
256
+ # ── Save dataset ────────────────────────────────────────
257
+ def save_dataset(records, nla_meta: dict):
258
+ """Save dataset as JSONL + stats."""
259
+ # Strip bulky activation vectors for JSONL (keep as separate array file)
260
+ jsonl_path = ARTIFACTS_DIR / "dataset.jsonl"
261
+ # Save activations as a separate .pt file for efficiency
262
+ act_path = ARTIFACTS_DIR / "activations.pt"
263
+
264
+ activation_tensors = []
265
+ json_records = []
266
+
267
+ for r in records:
268
+ act = r.pop("activation_vector")
269
+ activation_tensors.append(torch.tensor(act))
270
+ json_records.append(r)
271
+
272
+ with open(jsonl_path, "w", encoding="utf-8") as f:
273
+ for rec in json_records:
274
+ f.write(json.dumps(rec, ensure_ascii=False) + "\n")
275
+
276
+ torch.save(torch.stack(activation_tensors), act_path)
277
+
278
+ print(f"\n 💾 Dataset saved:")
279
+ print(f" {jsonl_path} ({len(json_records)} records)")
280
+ print(f" {act_path} ({activation_tensors[0].shape})")
281
+
282
+ # Stats
283
+ n_texts = len(set(r["text_idx"] for r in json_records))
284
+ has_explanation = sum(1 for r in json_records if r.get("teacher_explanation"))
285
+ empty_explanation = sum(1 for r in json_records if r.get("teacher_explanation") in ("[空输出]", None, ""))
286
+
287
+ stats = {
288
+ "total_records": len(json_records),
289
+ "texts_used": n_texts,
290
+ "records_with_explanation": has_explanation,
291
+ "records_empty_explanation": empty_explanation,
292
+ "d_model": activation_tensors[0].shape[0],
293
+ }
294
+
295
+ stats_path = ARTIFACTS_DIR / "dataset_stats.json"
296
+ with open(stats_path, "w", encoding="utf-8") as f:
297
+ json.dump(stats, f, ensure_ascii=False, indent=2)
298
+ print(f" {stats_path}")
299
+
300
+ # Update nla_meta
301
+ nla_meta["training"]["dataset_size"] = len(json_records)
302
+ with open(SIDECAR_PATH, "w", encoding="utf-8") as f:
303
+ import yaml
304
+ yaml.dump(nla_meta, f, allow_unicode=True, sort_keys=False, default_flow_style=False)
305
+ print(f" sidecar updated with dataset_size={len(json_records)}")
306
+
307
+ return json_records
308
+
309
+
310
+ # ── main ───────────────────────────────────────────────
311
+ def main():
312
+ print("=" * 60)
313
+ print("📦 TINY-NLA DATA GENERATION")
314
+ print("=" * 60)
315
+ print()
316
+
317
+ device = detect_device()
318
+ dtype = get_dtype(device)
319
+ print(f" Device: {device}, dtype: {dtype}")
320
+
321
+ # Load sidecar
322
+ with open(SIDECAR_PATH, "r") as f:
323
+ import yaml
324
+ nla_meta = yaml.safe_load(f)
325
+
326
+ layer_idx = nla_meta["layer_index"]
327
+ max_samples = 500
328
+
329
+ # ── Pass 1: Extract activations ──
330
+ print(f"\n Loading BASE model for activation extraction...")
331
+ t0 = time.perf_counter()
332
+ base_model = AutoModelForCausalLM.from_pretrained(
333
+ BASE_MODEL,
334
+ trust_remote_code=True,
335
+ torch_dtype=dtype,
336
+ low_cpu_mem_usage=True,
337
+ attn_implementation="eager",
338
+ ).to(device)
339
+ base_model.eval()
340
+ base_tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
341
+ print(f" Base model loaded in {time.perf_counter() - t0:.1f}s")
342
+
343
+ all_records = extract_activations(base_model, base_tokenizer, TEXTS, layer_idx, device, dtype)
344
+
345
+ # Clean up base model to free memory
346
+ del base_model
347
+ if device.type == "mps":
348
+ torch.mps.empty_cache()
349
+ print(" Base model freed from memory\n")
350
+
351
+ # ── Pass 2: Generate teacher explanations ──
352
+ print(f" Loading INSTRUCT model for teacher...")
353
+ t0 = time.perf_counter()
354
+ instruct_model = AutoModelForCausalLM.from_pretrained(
355
+ INSTRUCT_MODEL,
356
+ trust_remote_code=True,
357
+ torch_dtype=dtype,
358
+ low_cpu_mem_usage=True,
359
+ attn_implementation="eager",
360
+ ).to(device)
361
+ instruct_model.eval()
362
+ instruct_tokenizer = AutoTokenizer.from_pretrained(INSTRUCT_MODEL, trust_remote_code=True)
363
+ print(f" Instruct model loaded in {time.perf_counter() - t0:.1f}s")
364
+
365
+ results = generate_teacher_explanations(
366
+ instruct_model, instruct_tokenizer, all_records,
367
+ device, dtype, max_samples=max_samples,
368
+ )
369
+
370
+ # Clean up
371
+ del instruct_model
372
+ if device.type == "mps":
373
+ torch.mps.empty_cache()
374
+
375
+ # ── Save ──
376
+ json_records = save_dataset(results, nla_meta)
377
+
378
+ # ── Quality check ──
379
+ print()
380
+ print("=" * 60)
381
+ print("🔍 Quick Quality Check")
382
+ print("=" * 60)
383
+ sample = random.Random(42).sample(json_records, min(10, len(json_records)))
384
+ for r in sample:
385
+ exp = r.get("teacher_explanation", "")
386
+ print(f" [{r['token_text']!r:10}] {exp[:80]}")
387
+
388
+ print()
389
+ print("✅ Data generation complete!")
390
+ print(f" Dataset: {ARTIFACTS_DIR}/dataset.jsonl")
391
+ print(f" Activations: {ARTIFACTS_DIR}/activations.pt")
392
+ print(f" Stats: {ARTIFACTS_DIR}/dataset_stats.json")
393
+ print(f" Sidecar: {SIDECAR_PATH}")
394
+
395
+
396
+ if __name__ == "__main__":
397
+ main()
experiments/tiny_nla/generate_teacher_labels.py ADDED
@@ -0,0 +1,178 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Re-generate teacher explanations with correct parameters.
4
+ Qwen3-0.6B instruct uses <think> blocks; we:
5
+ 1. Use max_new_tokens=200 (enough for CoT + answer)
6
+ 2. Add explicit instruction to suppress CoT
7
+ 3. Post-process to strip any remaining <think> blocks
8
+ """
9
+
10
+ import json, re, os, time, random
11
+ from pathlib import Path
12
+ import torch
13
+ from transformers import AutoModelForCausalLM, AutoTokenizer
14
+
15
+ ARTIFACTS_DIR = Path(__file__).resolve().parents[2] / "artifacts" / "tiny_nla"
16
+ INSTRUCT_MODEL = "Qwen/Qwen3-0.6B"
17
+
18
+ def detect_device():
19
+ if torch.cuda.is_available():
20
+ return torch.device("cuda")
21
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
22
+ return torch.device("mps")
23
+ return torch.device("cpu")
24
+
25
+ def clean_explanation(text: str) -> str:
26
+ """Strip <think> blocks and normalize."""
27
+ if not text:
28
+ return ""
29
+ cleaned = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
30
+ cleaned = cleaned.strip()
31
+ # Remove any remaining XML-like tags
32
+ cleaned = re.sub(r'<[^>]+>', '', cleaned).strip()
33
+ return cleaned
34
+
35
+ def generate_explanations(records, batch_size=1, max_new_tokens=200):
36
+ device = detect_device()
37
+ dtype = torch.float32
38
+
39
+ print(f" Device: {device}, dtype: {dtype}")
40
+ print(f" Generating {len(records)} explanations with max_new_tokens={max_new_tokens}")
41
+
42
+ # Load instruct model
43
+ print(" Loading Qwen3-0.6B instruct...")
44
+ t0 = time.perf_counter()
45
+ model = AutoModelForCausalLM.from_pretrained(
46
+ INSTRUCT_MODEL,
47
+ trust_remote_code=True,
48
+ torch_dtype=dtype,
49
+ low_cpu_mem_usage=True,
50
+ attn_implementation="eager",
51
+ ).to(device)
52
+ model.eval()
53
+ tokenizer = AutoTokenizer.from_pretrained(INSTRUCT_MODEL, trust_remote_code=True)
54
+ print(f" Loaded in {time.perf_counter() - t0:.1f}s\n")
55
+
56
+ # Strong instruction to avoid reasoning blocks
57
+ system_msg = "你是一个模型可解释性专家。用1-2句中文简短回答,不要分析过程,不要推理步骤,直接给出解释。"
58
+
59
+ results = []
60
+ failures = 0
61
+
62
+ for i, rec in enumerate(records):
63
+ text = rec.get("text", "")
64
+ pos = rec.get("pos", 0)
65
+ token_text = rec.get("token_text", "")
66
+ top_tokens = rec.get("top_tokens", [])
67
+
68
+ user_msg = (
69
+ f"文本:{text}\n"
70
+ f"位置:第{pos}个token,该token文本是「{token_text}」\n"
71
+ f"预测的下一个token候选:{', '.join(top_tokens[:5])}\n\n"
72
+ f"用1-2句中文解释模型在该位置关注什么语义信息。直接回答,不要思考过程。"
73
+ )
74
+
75
+ messages = [
76
+ {"role": "system", "content": system_msg},
77
+ {"role": "user", "content": user_msg},
78
+ ]
79
+
80
+ prompt = tokenizer.apply_chat_template(
81
+ messages, tokenize=False, add_generation_prompt=True
82
+ )
83
+
84
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
85
+
86
+ raw_explanation = ""
87
+ try:
88
+ with torch.no_grad():
89
+ output_ids = model.generate(
90
+ **inputs,
91
+ max_new_tokens=max_new_tokens,
92
+ do_sample=False, # greedy for speed
93
+ pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
94
+ )
95
+
96
+ generated = output_ids[0][inputs["input_ids"].shape[1]:]
97
+ raw_explanation = tokenizer.decode(generated, skip_special_tokens=True).strip()
98
+
99
+ explanation = clean_explanation(raw_explanation)
100
+ if not explanation or len(explanation) < 5:
101
+ explanation = "[空输出]"
102
+ failures += 1
103
+
104
+ except Exception as e:
105
+ explanation = f"[生成失败: {e}]"
106
+ failures += 1
107
+
108
+ rec["teacher_explanation_raw"] = raw_explanation
109
+ rec["teacher_explanation"] = explanation
110
+
111
+ if (i + 1) % 20 == 0:
112
+ print(f" [{i+1}/{len(records)}] failures={failures}", flush=True)
113
+ if i < 3:
114
+ print(f" [{i+1}] {token_text!r:10} -> {explanation[:80]}")
115
+
116
+ print(f"\n ✅ Generated {len(records)} explanations, failures={failures}")
117
+
118
+ # Clean up
119
+ del model
120
+ if device.type == "mps":
121
+ torch.mps.empty_cache()
122
+
123
+ return records
124
+
125
+
126
+ def main():
127
+ print("=" * 60)
128
+ print("🔄 Re-Generating Teacher Explanations")
129
+ print("=" * 60)
130
+
131
+ # Load existing dataset
132
+ jsonl_path = ARTIFACTS_DIR / "dataset.jsonl"
133
+ records = []
134
+ with open(jsonl_path, "r", encoding="utf-8") as f:
135
+ for line in f:
136
+ records.append(json.loads(line))
137
+ print(f" Loaded {len(records)} records from {jsonl_path}")
138
+
139
+ # Regenerate explanations
140
+ records = generate_explanations(records, max_new_tokens=200)
141
+
142
+ # Quality stats
143
+ valid = [r for r in records if r.get("teacher_explanation") and r["teacher_explanation"] not in ("[空输出]", "")]
144
+ empty = len(records) - len(valid)
145
+ has_think = sum(1 for r in records if "<think>" in (r.get("teacher_explanation", "") or ""))
146
+
147
+ print(f"\n Quality: {len(valid)} valid, {empty} empty, {has_think} still has think tag")
148
+
149
+ # Save
150
+ with open(jsonl_path, "w", encoding="utf-8") as f:
151
+ for r in records:
152
+ f.write(json.dumps(r, ensure_ascii=False) + "\n")
153
+ print(f" ✅ Saved to {jsonl_path}")
154
+
155
+ # Save AV training data
156
+ av_data = [{
157
+ "text": r["text"],
158
+ "token_text": r["token_text"],
159
+ "pos": r["pos"],
160
+ "teacher_explanation": r["teacher_explanation"],
161
+ "top_tokens": r["top_tokens"],
162
+ } for r in records if r.get("teacher_explanation") and r["teacher_explanation"] not in ("[空输出]", "")]
163
+
164
+ av_path = ARTIFACTS_DIR / "av_training_data.json"
165
+ with open(av_path, "w", encoding="utf-8") as f:
166
+ json.dump(av_data, f, ensure_ascii=False, indent=2)
167
+ print(f" ✅ AV training data: {av_path} ({len(av_data)} samples)")
168
+
169
+ # Sample
170
+ print("\n Samples:")
171
+ for r in random.Random(42).sample(av_data, min(5, len(av_data))):
172
+ print(f" [{r['token_text']!r:10}] {r['teacher_explanation'][:80]}")
173
+
174
+ print("\n✅ Done!")
175
+
176
+
177
+ if __name__ == "__main__":
178
+ main()
experiments/tiny_nla/generate_teacher_labels_opencode.py ADDED
@@ -0,0 +1,207 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Generate high-quality teacher labels via opencode CLI (or any LLM).
4
+
5
+ Usage:
6
+ python generate_teacher_labels_opencode.py
7
+
8
+ This script:
9
+ 1. Reads dataset.jsonl (activation records with context/token/top_tokens)
10
+ 2. For each record, prints a prompt for the LLM (opencode/Claude/etc.)
11
+ 3. Reads LLM output and saves to teacher_labels_hq.json
12
+
13
+ The output file can then be merged into av_training_data.json for retraining.
14
+
15
+ Run mode options:
16
+ --dry-run Print prompts only (for review), no LLM calls
17
+ --local Use local Qwen3-0.6B instruct as fallback teacher
18
+ --opencode-cmd Path/name of opencode CLI (default: "opencode")
19
+ """
20
+
21
+ import json, sys, subprocess, argparse, time
22
+ from pathlib import Path
23
+
24
+ REPO_ROOT = Path(__file__).resolve().parents[2]
25
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
26
+ DATASET = ARTIFACTS_DIR / "dataset.jsonl"
27
+ OUTPUT = ARTIFACTS_DIR / "teacher_labels_hq.json"
28
+
29
+
30
+ TEACHER_PROMPT = """\
31
+ 你是一个语言模型内部机制分析专家。我需要你用 1-2 句简洁中文,描述一个 Transformer 语言模型(Qwen3-0.6B)在处理特定 token 时,该位置的 residual stream activation 所编码的语义信息。
32
+
33
+ ---
34
+ 上下文句子:{context}
35
+ 当前分析的 token:「{token}」(位置 {pos},共 {seq_len} 个 token)
36
+ 模型在此位置预测的下一个词(概率最高的候选):{top_tokens}
37
+ ---
38
+
39
+ 请根据上面信息,用 1-2 句中文描述:
40
+ - 模型在这个 token 的位置正在编码什么语义信息
41
+ - 这个位置的激活值如何帮助预测后续内容
42
+
43
+ 要求:
44
+ - 简洁具体,不要泛泛而谈
45
+ - 联系「{token}」在句子中的实际作用
46
+ - 联系模型预测的下一个词来推断编码内容
47
+ - 不要以"这个位置"或"激活值"开头,而是直接描述语义
48
+ - 不超过 60 字
49
+
50
+ 只输出解释本身,不要输出分析过程。"""
51
+
52
+
53
+ def make_prompt(rec):
54
+ top5 = "、".join(rec["top_tokens"][:5])
55
+ seq_len = len(rec["text"]) # approximate
56
+ return TEACHER_PROMPT.format(
57
+ context=rec["text"],
58
+ token=rec["token_text"],
59
+ pos=rec["pos"],
60
+ seq_len=seq_len,
61
+ top_tokens=top5,
62
+ )
63
+
64
+
65
+ def call_opencode(prompt_text, opencode_cmd="opencode", model="opencode/deepseek-v4-flash-free"):
66
+ """Call opencode CLI with a prompt, return response text."""
67
+ result = subprocess.run(
68
+ [opencode_cmd, "run", "--model", model, prompt_text],
69
+ capture_output=True,
70
+ text=True,
71
+ timeout=120,
72
+ )
73
+ if result.returncode != 0:
74
+ raise RuntimeError(f"opencode error: {result.stderr[:200]}")
75
+ # Strip ANSI escape codes and header lines
76
+ output = result.stdout
77
+ lines = output.splitlines()
78
+ # Skip lines starting with ANSI/control chars or "> orchestrator"
79
+ content_lines = [
80
+ l for l in lines
81
+ if l.strip() and not l.strip().startswith("\x1b") and not l.strip().startswith("> orchestrator")
82
+ ]
83
+ return "\n".join(content_lines).strip()
84
+
85
+
86
+ def call_local_qwen(prompt_text):
87
+ """Fallback: use local Qwen3-0.6B instruct."""
88
+ import torch
89
+ from transformers import AutoModelForCausalLM, AutoTokenizer
90
+
91
+ if not hasattr(call_local_qwen, "_model"):
92
+ print(" Loading local Qwen3-0.6B instruct...")
93
+ tok = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B", trust_remote_code=True)
94
+ mdl = AutoModelForCausalLM.from_pretrained(
95
+ "Qwen/Qwen3-0.6B", trust_remote_code=True, dtype=torch.float32,
96
+ low_cpu_mem_usage=True, attn_implementation="eager",
97
+ )
98
+ mdl.eval()
99
+ call_local_qwen._tok = tok
100
+ call_local_qwen._model = mdl
101
+
102
+ tok, mdl = call_local_qwen._tok, call_local_qwen._model
103
+ msgs = [{"role": "user", "content": prompt_text}]
104
+ text = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
105
+ inp = tok(text, return_tensors="pt")
106
+ with torch.no_grad():
107
+ out = mdl.generate(
108
+ **inp, max_new_tokens=128, do_sample=False,
109
+ pad_token_id=tok.eos_token_id, eos_token_id=tok.eos_token_id,
110
+ )
111
+ gen = out[0][inp["input_ids"].shape[1]:]
112
+ return tok.decode(gen, skip_special_tokens=True).strip()
113
+
114
+
115
+ def main():
116
+ parser = argparse.ArgumentParser()
117
+ parser.add_argument("--dry-run", action="store_true", help="Print prompts, no LLM calls")
118
+ parser.add_argument("--local", action="store_true", help="Use local Qwen instruct")
119
+ parser.add_argument("--opencode-cmd", default="opencode", help="opencode CLI command")
120
+ parser.add_argument("--model", default="opencode/deepseek-v4-flash-free", help="opencode model")
121
+ parser.add_argument("--limit", type=int, default=None, help="Process only first N records")
122
+ parser.add_argument("--skip-existing", action="store_true", default=True, help="Skip if output exists")
123
+ args = parser.parse_args()
124
+
125
+ with open(DATASET) as f:
126
+ records = [json.loads(l) for l in f]
127
+
128
+ if args.limit:
129
+ records = records[:args.limit]
130
+
131
+ # Filter out OOD first-token records (norm > 2000) — they degrade training
132
+ records = [r for r in records if r["activation_norm"] < 2000]
133
+ print(f"Filtered to {len(records)} in-distribution records (norm < 2000)")
134
+
135
+ # Load existing results if any
136
+ existing = {}
137
+ if OUTPUT.exists() and args.skip_existing:
138
+ with open(OUTPUT) as f:
139
+ existing_list = json.load(f)
140
+ existing = {(r["text_idx"], r["pos"]): r for r in existing_list}
141
+ print(f"Loaded {len(existing)} existing labels")
142
+
143
+ results = list(existing.values())
144
+ new_count = 0
145
+
146
+ for i, rec in enumerate(records):
147
+ key = (rec["text_idx"], rec["pos"])
148
+ if key in existing:
149
+ continue
150
+
151
+ prompt = make_prompt(rec)
152
+
153
+ if args.dry_run:
154
+ print(f"\n=== Record {i}: [{rec['token_text']}] pos={rec['pos']} ===")
155
+ print(prompt)
156
+ print("---")
157
+ continue
158
+
159
+ try:
160
+ if args.local:
161
+ explanation = call_local_qwen(prompt)
162
+ else:
163
+ explanation = call_opencode(prompt, args.opencode_cmd, args.model)
164
+
165
+ # Clean up: strip think blocks if any
166
+ if "<think>" in explanation:
167
+ if "</think>" in explanation:
168
+ explanation = explanation.split("</think>")[-1].strip()
169
+ else:
170
+ explanation = explanation.split("<think>")[0].strip()
171
+
172
+ result = {
173
+ "text_idx": rec["text_idx"],
174
+ "pos": rec["pos"],
175
+ "text": rec["text"],
176
+ "token_text": rec["token_text"],
177
+ "top_tokens": rec["top_tokens"][:5],
178
+ "activation_norm": rec["activation_norm"],
179
+ "teacher_explanation": explanation,
180
+ "teacher_source": "local_qwen" if args.local else "opencode",
181
+ }
182
+ results.append(result)
183
+ new_count += 1
184
+
185
+ if new_count % 10 == 0 or i == 0:
186
+ # Save checkpoint
187
+ with open(OUTPUT, "w", encoding="utf-8") as f:
188
+ json.dump(results, f, ensure_ascii=False, indent=2)
189
+ print(f" [{i+1}/{len(records)}] saved {len(results)} labels")
190
+ print(f" [{rec['token_text']}] → {explanation[:60]}")
191
+
192
+ time.sleep(0.3) # gentle rate limit
193
+
194
+ except Exception as e:
195
+ print(f" Error on record {i}: {e}", file=sys.stderr)
196
+ time.sleep(2)
197
+
198
+ if not args.dry_run:
199
+ with open(OUTPUT, "w", encoding="utf-8") as f:
200
+ json.dump(results, f, ensure_ascii=False, indent=2)
201
+ print(f"\nDone. {len(results)} total labels saved to {OUTPUT}")
202
+ else:
203
+ print(f"\nDry run complete. {len(records)} prompts shown.")
204
+
205
+
206
+ if __name__ == "__main__":
207
+ main()
experiments/tiny_nla/generate_teacher_labels_v2.py ADDED
@@ -0,0 +1,186 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Three-channel parallel teacher label generation.
4
+ Output: append-mode JSONL (one record per line) - no corruption possible.
5
+ Channels: NVIDIA API + OpenRouter (gpt-oss-120b, nex-n2-pro) + opencode CLI
6
+ Usage: python generate_teacher_labels_v2.py [--workers N]
7
+ """
8
+ import json, os, subprocess, time, argparse
9
+ from pathlib import Path
10
+ from concurrent.futures import ThreadPoolExecutor, as_completed
11
+ from openai import OpenAI
12
+
13
+ ARTIFACTS = Path(__file__).resolve().parents[2] / "artifacts" / "tiny_nla"
14
+ RECORDS_FILE = ARTIFACTS / "records_v2.jsonl"
15
+ OUTPUT_JSONL = ARTIFACTS / "teacher_labels_v2.jsonl" # append-mode, safe for concurrent writes
16
+ OUTPUT_JSON = ARTIFACTS / "teacher_labels_v2.json" # final merged output
17
+
18
+ NVIDIA_KEY = os.environ.get("NVIDIA_API_KEY", "")
19
+ OR_KEY = os.environ.get("OPENROUTER_API_KEY", "")
20
+ DS_KEY = os.environ.get("DEEPSEEK_API_KEY", "")
21
+ FW_KEY = os.environ.get("FIREWORKS_API_KEY", "")
22
+
23
+ nvidia_client = OpenAI(base_url="https://integrate.api.nvidia.com/v1", api_key=NVIDIA_KEY)
24
+ or_client = OpenAI(base_url="https://openrouter.ai/api/v1", api_key=OR_KEY)
25
+ ds_client = OpenAI(base_url="https://api.deepseek.com", api_key=DS_KEY)
26
+ fw_client = OpenAI(base_url="https://api.fireworks.ai/inference/v1", api_key=FW_KEY)
27
+
28
+ PROMPT = """\
29
+ 你是语言模型内部机制分析专家。用1-2句简洁中文描述Qwen3-0.6B在处理下面这个token时,该位置residual stream激活值编码的语义信息。
30
+
31
+ 上下文:{context}
32
+ 当前token:「{token}」(位置{pos}/{seq_len})
33
+ 模型预测下一个词的候选:{top_tokens}
34
+
35
+ 要求:简洁具体,联系token在句中的实际句法/语义角色,联系预测候选推断编码内容,不超过55字。只输出解释本身。"""
36
+
37
+
38
+ def make_prompt(r):
39
+ return PROMPT.format(context=r["text"], token=r["token_text"],
40
+ pos=r["pos"], seq_len=r["seq_len"],
41
+ top_tokens="、".join(r["top_tokens"][:5]))
42
+
43
+
44
+ def strip_think(t):
45
+ return t.split("</think>")[-1].strip() if "</think>" in t else t.strip()
46
+
47
+
48
+ def call_nvidia(r):
49
+ resp = nvidia_client.chat.completions.create(
50
+ model="deepseek-ai/deepseek-v4-pro",
51
+ messages=[{"role": "user", "content": make_prompt(r)}], max_tokens=256)
52
+ return strip_think(resp.choices[0].message.content or ""), "nvidia"
53
+
54
+
55
+ def call_gpt_oss(r):
56
+ resp = or_client.chat.completions.create(
57
+ model="openai/gpt-oss-120b:free",
58
+ messages=[{"role": "user", "content": make_prompt(r)}], max_tokens=256)
59
+ return strip_think(resp.choices[0].message.content or ""), "or-gpt-oss"
60
+
61
+
62
+ def call_nex(r):
63
+ resp = or_client.chat.completions.create(
64
+ model="nex-agi/nex-n2-pro:free",
65
+ messages=[{"role": "user", "content": make_prompt(r)}], max_tokens=256)
66
+ return strip_think(resp.choices[0].message.content or ""), "or-nex"
67
+
68
+
69
+ def call_deepseek(r):
70
+ resp = ds_client.chat.completions.create(
71
+ model="deepseek-chat",
72
+ messages=[{"role": "user", "content": make_prompt(r)}], max_tokens=256)
73
+ return strip_think(resp.choices[0].message.content or ""), "deepseek"
74
+
75
+
76
+ def call_fireworks(r):
77
+ resp = fw_client.chat.completions.create(
78
+ model="accounts/fireworks/models/gpt-oss-120b",
79
+ messages=[{"role": "user", "content": make_prompt(r)}], max_tokens=256)
80
+ return strip_think(resp.choices[0].message.content or ""), "fireworks"
81
+
82
+
83
+ # DeepSeek + Fireworks as primary (both confirmed working), others as fallback
84
+ CHANNELS = [call_deepseek, call_fireworks, call_deepseek, call_fireworks,
85
+ call_deepseek, call_fireworks, call_nvidia, call_gpt_oss]
86
+
87
+
88
+ def process(args):
89
+ idx, rec, ch_idx = args
90
+ order = [CHANNELS[(ch_idx + i) % len(CHANNELS)] for i in range(len(CHANNELS))]
91
+ for fn in order:
92
+ try:
93
+ expl, src = fn(rec)
94
+ if expl and len(expl) >= 5:
95
+ return {**rec, "teacher_explanation": expl, "teacher_source": src}
96
+ except Exception:
97
+ continue
98
+ return None
99
+
100
+
101
+ def load_done():
102
+ """Load all completed keys from existing JSON and JSONL."""
103
+ done = {}
104
+ if OUTPUT_JSON.exists():
105
+ try:
106
+ for r in json.load(open(OUTPUT_JSON)):
107
+ done[(r["text_idx"], r["pos"])] = True
108
+ except Exception:
109
+ pass
110
+ if OUTPUT_JSONL.exists():
111
+ for line in open(OUTPUT_JSONL):
112
+ try:
113
+ r = json.loads(line)
114
+ done[(r["text_idx"], r["pos"])] = True
115
+ except Exception:
116
+ pass
117
+ return done
118
+
119
+
120
+ def main():
121
+ parser = argparse.ArgumentParser()
122
+ parser.add_argument("--workers", type=int, default=16)
123
+ args = parser.parse_args()
124
+
125
+ records = [json.loads(l) for l in open(RECORDS_FILE)]
126
+ done = load_done()
127
+ todo = [(i, r, i) for i, r in enumerate(records)
128
+ if (r["text_idx"], r["pos"]) not in done]
129
+
130
+ print(f"Records: {len(records)}, done: {len(done)}, todo: {len(todo)}, workers: {args.workers}")
131
+ print(f"Channels: nvidia / gpt-oss / nex-n2-pro / opencode (round-robin)")
132
+
133
+ completed = 0
134
+ errors = 0
135
+ t0 = time.time()
136
+
137
+ # Open jsonl in append mode - each write is atomic (small line, OS guarantees)
138
+ out_f = open(OUTPUT_JSONL, "a", encoding="utf-8", buffering=1) # line-buffered
139
+
140
+ with ThreadPoolExecutor(max_workers=args.workers) as pool:
141
+ futures = {pool.submit(process, item): item for item in todo}
142
+ for fut in as_completed(futures):
143
+ result = fut.result()
144
+ if result:
145
+ out_f.write(json.dumps(result, ensure_ascii=False) + "\n")
146
+ completed += 1
147
+ else:
148
+ errors += 1
149
+
150
+ n = completed + errors
151
+ if n % 200 == 0:
152
+ rate = n / (time.time() - t0) * 60
153
+ remaining = len(todo) - n
154
+ eta_min = remaining / (rate / 60) if rate > 0 else 0
155
+ print(f" [{len(done)+completed}/{len(records)}] errors={errors} "
156
+ f"rate={rate:.0f}/min ETA={eta_min/60:.1f}h")
157
+
158
+ out_f.close()
159
+
160
+ # Merge jsonl → final json
161
+ all_done = {}
162
+ if OUTPUT_JSON.exists():
163
+ try:
164
+ for r in json.load(open(OUTPUT_JSON)):
165
+ all_done[(r["text_idx"], r["pos"])] = r
166
+ except Exception:
167
+ pass
168
+ for line in open(OUTPUT_JSONL):
169
+ try:
170
+ r = json.loads(line)
171
+ all_done[(r["text_idx"], r["pos"])] = r
172
+ except Exception:
173
+ pass
174
+ with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
175
+ json.dump(list(all_done.values()), f, ensure_ascii=False, indent=2)
176
+
177
+ print(f"\nDone: {len(all_done)} labels → {OUTPUT_JSON}")
178
+
179
+ # Sample
180
+ items = list(all_done.values())[-5:]
181
+ for item in items:
182
+ print(f" [{item['teacher_source']}] {item['token_text']} → {item['teacher_explanation'][:65]}")
183
+
184
+
185
+ if __name__ == "__main__":
186
+ main()
experiments/tiny_nla/infer_tiny_nla.py ADDED
@@ -0,0 +1,316 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Tiny-NLA Inference Script
4
+
5
+ Reads nla_meta.yaml and provides:
6
+ extract(text, positions) → activation vectors at layer 19
7
+ explain(activation) → AV-generated explanation
8
+ reconstruct(explanation) → AR-reconstructed activation + cosine/MSE
9
+
10
+ Usage:
11
+ python infer_tiny_nla.py --text "你的文本" --position 5
12
+ python infer_tiny_nla.py --interactive
13
+ """
14
+
15
+ import argparse, json, os, sys, yaml, time
16
+ from pathlib import Path
17
+
18
+ import torch
19
+ import torch.nn.functional as F
20
+ from transformers import AutoModelForCausalLM, AutoTokenizer
21
+ from peft import PeftModel
22
+
23
+ # ── paths ──────────────────────────────────────────────
24
+ REPO_ROOT = Path(__file__).resolve().parents[2]
25
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
26
+ CHECKPOINT_DIR = REPO_ROOT / "artifacts" / "tiny_nla" / "checkpoints"
27
+ AV_CHECKPOINT = CHECKPOINT_DIR / "av"
28
+ AR_CHECKPOINT = CHECKPOINT_DIR / "ar" / "best_ar_head.pt"
29
+
30
+
31
+ class TinyNLA:
32
+ """
33
+ Tiny-NLA inference wrapper.
34
+ Loads models lazily to avoid memory waste.
35
+ """
36
+
37
+ def __init__(self, sidecar_path=SIDECAR_PATH):
38
+ with open(sidecar_path, "r") as f:
39
+ self.meta = yaml.safe_load(f)
40
+
41
+ self.base_model_name = self.meta["base_model"]
42
+ self.av_model_name = self.meta["av_init_model"]
43
+ self.layer_idx = self.meta["layer_index"]
44
+ self.d_model = self.meta["d_model"]
45
+ self.inj_char = self.meta["tokens"]["injection_char"]
46
+ self.inj_token_id = self.meta["tokens"]["injection_token_id"]
47
+ self.inj_scale = self.meta["extraction"]["injection_scale"]
48
+
49
+ self.device = self._detect_device()
50
+ self.dtype = torch.float32
51
+
52
+ # Lazy-loaded models
53
+ self.base_model = None
54
+ self.base_tokenizer = None
55
+ self.av_model = None
56
+ self.av_tokenizer = None
57
+ self.ar_head = None
58
+ self.ar_loaded = False
59
+
60
+ def _detect_device(self):
61
+ if torch.cuda.is_available():
62
+ return torch.device("cuda")
63
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
64
+ return torch.device("mps")
65
+ return torch.device("cpu")
66
+
67
+ def _load_base(self):
68
+ if self.base_model is not None:
69
+ return
70
+ print(f" Loading base model ({self.base_model_name})...")
71
+ t0 = time.perf_counter()
72
+ self.base_model = AutoModelForCausalLM.from_pretrained(
73
+ self.base_model_name,
74
+ trust_remote_code=True,
75
+ torch_dtype=self.dtype,
76
+ low_cpu_mem_usage=True,
77
+ attn_implementation="eager",
78
+ ).to(self.device)
79
+ self.base_model.eval()
80
+ self.base_tokenizer = AutoTokenizer.from_pretrained(
81
+ self.base_model_name, trust_remote_code=True
82
+ )
83
+ print(f" Done in {time.perf_counter() - t0:.1f}s")
84
+
85
+ def _load_av(self):
86
+ if self.av_model is not None:
87
+ return
88
+ print(f" Loading AV model ({self.av_model_name} + LoRA)...")
89
+ t0 = time.perf_counter()
90
+ base = AutoModelForCausalLM.from_pretrained(
91
+ self.av_model_name,
92
+ trust_remote_code=True,
93
+ torch_dtype=self.dtype,
94
+ low_cpu_mem_usage=True,
95
+ attn_implementation="eager",
96
+ ).to(self.device)
97
+ self.av_model = PeftModel.from_pretrained(base, AV_CHECKPOINT)
98
+ self.av_model.eval()
99
+ self.av_tokenizer = AutoTokenizer.from_pretrained(
100
+ self.av_model_name, trust_remote_code=True
101
+ )
102
+ print(f" Done in {time.perf_counter() - t0:.1f}s")
103
+
104
+ def _load_ar(self):
105
+ if self.ar_loaded:
106
+ return
107
+ if not AR_CHECKPOINT.exists():
108
+ print(f" ⚠️ AR checkpoint not found at {AR_CHECKPOINT}")
109
+ self.ar_loaded = False
110
+ return
111
+ print(f" Loading AR head...")
112
+ state = torch.load(AR_CHECKPOINT, map_location=self.device, weights_only=True)
113
+ # ARHead saved as {'linear.weight': ...} but we want {'weight': ...}
114
+ if "linear.weight" in state:
115
+ state = {"weight": state["linear.weight"]}
116
+ self.ar_head = torch.nn.Linear(self.d_model, self.d_model, bias=False)
117
+ self.ar_head.load_state_dict(state)
118
+ self.ar_head.to(self.device)
119
+ self.ar_head.eval()
120
+ self.ar_loaded = True
121
+ print(f" Done")
122
+
123
+ def extract(self, text: str, position: int = -1):
124
+ """
125
+ Extract activation at given layer for the specified token position.
126
+ If position < 0, uses last token.
127
+ Returns: (activation_vector [d_model], input_ids, top_tokens)
128
+ """
129
+ self._load_base()
130
+
131
+ inputs = self.base_tokenizer(text, return_tensors="pt").to(self.device)
132
+ seq_len = inputs["input_ids"].shape[1]
133
+
134
+ if position < 0 or position >= seq_len:
135
+ position = seq_len - 1
136
+
137
+ with torch.no_grad():
138
+ outputs = self.base_model(
139
+ **inputs,
140
+ output_hidden_states=True,
141
+ )
142
+
143
+ hidden = outputs.hidden_states[self.layer_idx] # [1, seq, d_model]
144
+ activation = hidden[0, position, :].cpu() # [d_model]
145
+
146
+ # Top-k from logits at this position
147
+ logits = outputs.logits[0, position, :]
148
+ topk_vals, topk_idxs = torch.topk(logits, k=10)
149
+ top_tokens = [self.base_tokenizer.decode([t]) for t in topk_idxs]
150
+
151
+ return {
152
+ "activation": activation,
153
+ "position": position,
154
+ "seq_len": seq_len,
155
+ "token_text": self.base_tokenizer.decode([inputs["input_ids"][0, position].item()]),
156
+ "top_tokens": top_tokens,
157
+ "activation_norm": activation.norm().item(),
158
+ }
159
+
160
+ def explain(self, activation: torch.Tensor, max_new_tokens=64):
161
+ """Generate AV explanation from activation vector."""
162
+ self._load_av()
163
+
164
+ prompt = f"<concept>{self.inj_char}</concept>\n<explanation>"
165
+ inputs = self.av_tokenizer(prompt, return_tensors="pt").to(self.device)
166
+
167
+ # Inject activation at injection token position
168
+ embeds = self.av_model.get_input_embeddings()(inputs["input_ids"])
169
+ inj_positions = (inputs["input_ids"][0] == self.inj_token_id).nonzero(as_tuple=True)[0]
170
+
171
+ if len(inj_positions) > 0:
172
+ inj_pos = inj_positions[0].item()
173
+ # Normalize activation to injection_scale norm
174
+ norm = activation.norm()
175
+ if norm > 0:
176
+ activation = activation / norm * self.inj_scale
177
+ scaled_act = activation.to(embeds.dtype).to(self.device)
178
+ embeds[0, inj_pos, :] = scaled_act
179
+
180
+ with torch.no_grad():
181
+ output_ids = self.av_model.generate(
182
+ inputs_embeds=embeds,
183
+ max_new_tokens=max_new_tokens,
184
+ do_sample=False,
185
+ pad_token_id=self.av_tokenizer.pad_token_id or self.av_tokenizer.eos_token_id,
186
+ )
187
+
188
+ # Extract just the generated part (token-level, not string-level)
189
+ prompt_len_tokens = inputs["input_ids"].shape[1]
190
+ gen_token_ids = output_ids[0][prompt_len_tokens:]
191
+ explanation = self.av_tokenizer.decode(gen_token_ids, skip_special_tokens=True).strip()
192
+
193
+ return {
194
+ "explanation": explanation,
195
+ }
196
+
197
+ def reconstruct(self, explanation: str):
198
+ """AR: explanation text → reconstructed activation + metrics."""
199
+ self._load_ar()
200
+ self._load_base()
201
+
202
+ if not self.ar_loaded:
203
+ return {"error": "AR checkpoint not available"}
204
+
205
+ # Tokenize explanation with base model's tokenizer
206
+ inputs = self.base_tokenizer(
207
+ explanation,
208
+ return_tensors="pt",
209
+ truncation=True,
210
+ max_length=128,
211
+ ).to(self.device)
212
+
213
+ with torch.no_grad():
214
+ outputs = self.base_model(
215
+ **inputs,
216
+ output_hidden_states=True,
217
+ )
218
+ last_hidden = outputs.hidden_states[-1]
219
+ seq_len = inputs["attention_mask"].sum(dim=1) - 1
220
+ last_token_hidden = last_hidden[0, seq_len[0], :]
221
+
222
+ reconstructed = self.ar_head(last_token_hidden)
223
+
224
+ return {"reconstructed": reconstructed.cpu()}
225
+
226
+ def roundtrip(self, text: str, position: int = -1):
227
+ """Full round-trip: extract → explain → reconstruct."""
228
+ # Extract
229
+ ext = self.extract(text, position)
230
+ activation = ext["activation"]
231
+
232
+ # Explain
233
+ expl = self.explain(activation)
234
+
235
+ # Reconstruct
236
+ rec = self.reconstruct(expl["explanation"])
237
+
238
+ result = {
239
+ "text": text,
240
+ "position": ext["position"],
241
+ "token_text": ext["token_text"],
242
+ "activation_norm": ext["activation_norm"],
243
+ "top_tokens": ext["top_tokens"],
244
+ "av_explanation": expl["explanation"],
245
+ }
246
+
247
+ if "reconstructed" in rec:
248
+ orig_n = F.normalize(activation.unsqueeze(0), dim=-1)
249
+ recon = rec["reconstructed"].unsqueeze(0)
250
+ recon_n = F.normalize(recon, dim=-1)
251
+ cosine = (orig_n * recon_n).sum(dim=-1).item()
252
+ mse = F.mse_loss(orig_n, recon_n).item()
253
+
254
+ result["ar_cosine"] = round(cosine, 4)
255
+ result["ar_normalized_mse"] = round(mse, 6)
256
+
257
+ return result
258
+
259
+
260
+ def main():
261
+ parser = argparse.ArgumentParser(description="Tiny-NLA Inference")
262
+ parser.add_argument("--text", type=str, help="Input text")
263
+ parser.add_argument("--position", type=int, default=-1, help="Token position (-1 = last)")
264
+ parser.add_argument("--interactive", action="store_true", help="Interactive mode")
265
+ parser.add_argument("--batch", type=str, help="JSON file with list of {text, position}")
266
+ args = parser.parse_args()
267
+
268
+ nla = TinyNLA()
269
+
270
+ if args.interactive:
271
+ print("Tiny-NLA Interactive Mode (Ctrl+D to exit)\n")
272
+ while True:
273
+ try:
274
+ text = input("Text: ").strip()
275
+ if not text:
276
+ continue
277
+ pos_input = input("Position (default=last): ").strip()
278
+ pos = int(pos_input) if pos_input else -1
279
+
280
+ print("\n Processing...")
281
+ result = nla.roundtrip(text, pos)
282
+
283
+ print(f"\n Token: {result['token_text']!r}")
284
+ print(f" Position: {result['position']}")
285
+ print(f" Top tokens: {result['top_tokens'][:5]}")
286
+ print(f" AV: {result['av_explanation'][:120]}")
287
+ if "ar_cosine" in result:
288
+ print(f" AR cosine: {result['ar_cosine']}")
289
+ print()
290
+ except EOFError:
291
+ break
292
+
293
+ elif args.text:
294
+ result = nla.roundtrip(args.text, args.position)
295
+ print(json.dumps(result, ensure_ascii=False, indent=2))
296
+
297
+ elif args.batch:
298
+ with open(args.batch, "r", encoding="utf-8") as f:
299
+ batch = json.load(f)
300
+ results = []
301
+ for item in batch:
302
+ r = nla.roundtrip(item["text"], item.get("position", -1))
303
+ results.append(r)
304
+ print(f" [{len(results)}/{len(batch)}] {r['token_text']!r} -> cosine={r.get('ar_cosine', 'N/A')}")
305
+
306
+ out_path = Path(args.batch).parent / "roundtrip_results.json"
307
+ with open(out_path, "w", encoding="utf-8") as f:
308
+ json.dump(results, f, ensure_ascii=False, indent=2)
309
+ print(f"\n Results saved to {out_path}")
310
+
311
+ else:
312
+ parser.print_help()
313
+
314
+
315
+ if __name__ == "__main__":
316
+ main()
experiments/tiny_nla/nla_meta.yaml ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ kind: tiny_nla_model
2
+ base_model: Qwen/Qwen3-0.6B-Base
3
+ av_init_model: Qwen/Qwen3-0.6B
4
+ layer_index: 19
5
+ num_hidden_layers: 28
6
+ d_model: 1024
7
+ activation_source: residual_stream
8
+ token_position_policy: selected_token
9
+ extraction:
10
+ injection_scale: 126.223
11
+ mse_normalization: l2_direction
12
+ tokens:
13
+ injection_char: ㈎
14
+ injection_token_id: 149705
15
+ prompt_templates:
16
+ av: '<concept>{injection_char}</concept>
17
+
18
+ <explanation>'
19
+ ar: <explanation>{explanation_text}</explanation>
20
+ training:
21
+ device: mps
22
+ dtype: torch.float32
23
+ dataset_size: 284
24
+ teacher: local_qwen_instruct
25
+ created_at: '2026-06-16T21:43:52'
26
+ av_results:
27
+ best_val_loss: 0.6400
28
+ lora_r: 8
29
+ lora_alpha: 16
30
+ checkpoint: artifacts/tiny_nla/checkpoints/av
31
+ ar_results:
32
+ best_val_cosine: 0.6631
33
+ best_val_loss: 0.000658
34
+ mean_baseline_cosine: 0.5962
35
+ shuffled_baseline_cosine: 0.5252
36
+ improvement_over_mean: 0.0669
37
+ checkpoint: artifacts/tiny_nla/checkpoints/ar/best_ar_head.pt
38
+ activation_stats:
39
+ mean: 125.85986328125
40
+ p50: 126.22295379638672
41
+ p90: 127.06012725830078
42
+ max: 129.3959503173828
43
+ min: 120.7603988647461
44
+ std: 3.65517520904541
45
+ smoke_test:
46
+ forward_ok: true
47
+ gen_ok: true
48
+ cosine_post_injection: 0.5979639887809753
experiments/tiny_nla/postprocess_dataset.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Post-process dataset: clean teacher explanations by stripping <think> blocks.
4
+ Qwen3 instruct outputs reasoning in <think>...</think> tags.
5
+ We want only the text after </think>.
6
+
7
+ Also generates an improved dataset with augmented prompt templates.
8
+ """
9
+
10
+ import json, re, random
11
+ from pathlib import Path
12
+
13
+ REPO_ROOT = Path(__file__).resolve().parents[2]
14
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
15
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
16
+
17
+ JSONL_PATH = ARTIFACTS_DIR / "dataset.jsonl"
18
+ ACT_PATH = ARTIFACTS_DIR / "activations.pt"
19
+
20
+
21
+ def clean_explanation(text: str) -> str:
22
+ """Strip <think> blocks, extract only the answer text."""
23
+ if not text:
24
+ return ""
25
+ # Remove <think>...</think> blocks (possibly with newlines)
26
+ cleaned = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL)
27
+ cleaned = cleaned.strip()
28
+ # Also strip leading/trailing whitespace per line
29
+ cleaned = '\n'.join(line.strip() for line in cleaned.split('\n') if line.strip())
30
+ # If after stripping think there's still nothing, keep minimal
31
+ if not cleaned:
32
+ # Try to salvage something from within think
33
+ inner = re.search(r'<think>(.*?)</think>', text, re.DOTALL)
34
+ if inner:
35
+ cleaned = inner.group(1).strip()[:100]
36
+ return cleaned
37
+
38
+
39
+ def is_empty_or_useless(text: str) -> bool:
40
+ """Check if explanation is empty or just boilerplate."""
41
+ if not text or len(text) < 5:
42
+ return True
43
+ useless = ["[空输出]", "[生成失败", "<think>"]
44
+ if any(u in text for u in useless):
45
+ return True
46
+ return False
47
+
48
+
49
+ def main():
50
+ print("=" * 60)
51
+ print("🧹 Post-Processing Dataset")
52
+ print("=" * 60)
53
+
54
+ # Load dataset
55
+ records = []
56
+ with open(JSONL_PATH, "r", encoding="utf-8") as f:
57
+ for line in f:
58
+ records.append(json.loads(line))
59
+ print(f" Loaded {len(records)} records from {JSONL_PATH}")
60
+
61
+ # Check stats before
62
+ empty_before = sum(1 for r in records if is_empty_or_useless(r.get("teacher_explanation", "")))
63
+ has_think = sum(1 for r in records if "<think>" in r.get("teacher_explanation", ""))
64
+ print(f" Before cleaning:")
65
+ print(f" Empty/useless: {empty_before}")
66
+ print(f" Contains <think>: {has_think}")
67
+
68
+ # Clean explanations
69
+ for r in records:
70
+ raw = r.get("teacher_explanation", "")
71
+ cleaned = clean_explanation(raw)
72
+ r["teacher_explanation_raw"] = raw # keep original
73
+ r["teacher_explanation"] = cleaned
74
+
75
+ empty_after = sum(1 for r in records if is_empty_or_useless(r.get("teacher_explanation", "")))
76
+ print(f" After cleaning:")
77
+ print(f" Empty/useless: {empty_after}")
78
+
79
+ # Check quality
80
+ print("\n Sample cleaned explanations:")
81
+ sample = random.Random(42).sample(records, 8)
82
+ for r in sample:
83
+ print(f" [{r['token_text']!r:10}] {r['teacher_explanation'][:80]}")
84
+
85
+ # Save cleaned dataset
86
+ with open(JSONL_PATH, "w", encoding="utf-8") as f:
87
+ for r in records:
88
+ f.write(json.dumps(r, ensure_ascii=False) + "\n")
89
+ print(f"\n ✅ Saved cleaned dataset to {JSONL_PATH}")
90
+
91
+ # Also save a separate "av ready" format (for AV training directly)
92
+ av_data = []
93
+ for r in records:
94
+ if not is_empty_or_useless(r.get("teacher_explanation", "")):
95
+ av_data.append({
96
+ "text": r["text"],
97
+ "token_text": r["token_text"],
98
+ "pos": r["pos"],
99
+ "teacher_explanation": r["teacher_explanation"],
100
+ "top_tokens": r["top_tokens"],
101
+ })
102
+
103
+ av_path = ARTIFACTS_DIR / "av_training_data.json"
104
+ with open(av_path, "w", encoding="utf-8") as f:
105
+ json.dump(av_data, f, ensure_ascii=False, indent=2)
106
+ print(f" ✅ Saved AV-ready data ({len(av_data)} samples) to {av_path}")
107
+
108
+ print("\n✅ Post-processing complete!")
109
+
110
+
111
+ if __name__ == "__main__":
112
+ main()
experiments/tiny_nla/smoke_stage0.py ADDED
@@ -0,0 +1,401 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Stage 0: Environment & Injection Smoke
4
+
5
+ Loads Qwen3-0.6B-Base, checks device, inspects model config,
6
+ extracts activations from ~2/3 depth layer, finds stats,
7
+ identifies a single-token injection character, and runs
8
+ an input_embeds injection smoke test.
9
+
10
+ Output: prints report + writes nla_meta.yaml sidecar
11
+ """
12
+
13
+ import os, sys, json, math, yaml, time
14
+ from pathlib import Path
15
+
16
+ import torch
17
+ import torch.nn.functional as F
18
+ from transformers import AutoModelForCausalLM, AutoTokenizer
19
+
20
+ # ── paths ──────────────────────────────────────────────
21
+ REPO_ROOT = Path(__file__).resolve().parents[2]
22
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
23
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
24
+ ARTIFACTS_DIR.mkdir(parents=True, exist_ok=True)
25
+
26
+ BASE_MODEL = "Qwen/Qwen3-0.6B-Base"
27
+ INSTRUCT_MODEL = "Qwen/Qwen3-0.6B"
28
+
29
+ # ── device detection ───────────────────────────────────
30
+ def detect_device():
31
+ """MPS preferred, CPU fallback. Prints diagnostics."""
32
+ print("=" * 60)
33
+ print("🔧 Device Detection")
34
+ print("=" * 60)
35
+ if torch.cuda.is_available():
36
+ device = torch.device("cuda")
37
+ print(f" ✅ CUDA available: {torch.cuda.get_device_name()}")
38
+ elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
39
+ device = torch.device("mps")
40
+ print(f" ✅ MPS available (Apple Silicon)")
41
+ else:
42
+ device = torch.device("cpu")
43
+ print(f" ⚠️ No accelerator, using CPU")
44
+ print(f" → Using device: {device}")
45
+ print()
46
+ return device
47
+
48
+
49
+ def mps_backward_smoke(device: torch.device):
50
+ """Quick backward pass test on MPS to detect flaky ops."""
51
+ if device.type != "mps":
52
+ return True
53
+ print(" 🔬 MPS backward smoke test...")
54
+ try:
55
+ x = torch.randn(2, 64, 1536, device=device, requires_grad=True)
56
+ loss = (x ** 2).mean()
57
+ loss.backward()
58
+ print(f" ✅ MPS backward: OK (grad norm={x.grad.norm():.4f})")
59
+ return True
60
+ except Exception as e:
61
+ print(f" ❌ MPS backward FAILED: {e}")
62
+ print(f" ⚠️ Falling back to CPU")
63
+ return False
64
+
65
+
66
+ def get_dtype_for_device(device: torch.device):
67
+ """Pick dtype per device. MPS float16 can be flaky → use float32."""
68
+ if device.type == "cuda":
69
+ return torch.float16
70
+ elif device.type == "mps":
71
+ # MPS float16 backward can fail on some ops; float32 is safe
72
+ print(" ℹ️ Using float32 on MPS (safer for backward)")
73
+ return torch.float32
74
+ return torch.float32
75
+
76
+
77
+ # ── model loading ──────────────────────────────────────
78
+ def load_model_and_tokenizer(model_name: str, device: torch.device, dtype: torch.dtype):
79
+ """Load model in eval mode on target device with proper settings."""
80
+ print(f" Loading {model_name}...")
81
+ t0 = time.perf_counter()
82
+ model = AutoModelForCausalLM.from_pretrained(
83
+ model_name,
84
+ trust_remote_code=True,
85
+ torch_dtype=dtype,
86
+ low_cpu_mem_usage=True,
87
+ attn_implementation="eager", # safest for MPS/CPU
88
+ ).to(device)
89
+ model.eval()
90
+ tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
91
+ load_time = time.perf_counter() - t0
92
+ n_params = sum(p.numel() for p in model.parameters())
93
+ print(f" ✅ Loaded: {n_params/1e6:.1f}M params in {load_time:.1f}s")
94
+ return model, tok
95
+
96
+
97
+ # ── model inspection ───────────────────────────────────
98
+ def inspect_model(model) -> dict:
99
+ """Read config and print architecture info."""
100
+ cfg = model.config
101
+ info = {
102
+ "num_hidden_layers": cfg.num_hidden_layers,
103
+ "hidden_size": cfg.hidden_size,
104
+ "vocab_size": cfg.vocab_size,
105
+ "intermediate_size": cfg.intermediate_size,
106
+ "num_attention_heads": cfg.num_attention_heads,
107
+ "num_key_value_heads": getattr(cfg, "num_key_value_heads", cfg.num_attention_heads),
108
+ }
109
+ info["layer_index"] = round(info["num_hidden_layers"] * 2 / 3)
110
+
111
+ print("=" * 60)
112
+ print("📐 Model Architecture")
113
+ print("=" * 60)
114
+ for k, v in info.items():
115
+ print(f" {k}: {v}")
116
+ print(f" layer_index (2/3 depth): {info['layer_index']}")
117
+ print()
118
+ return info
119
+
120
+
121
+ # ── activation extraction ──────────────────────────────
122
+ def extract_activation(model, tokenizer, text: str, layer_idx: int, device: torch.device):
123
+ """
124
+ Run forward with output_hidden_states=True, extract residual stream
125
+ at given layer for the last token position.
126
+ Returns: (activation_vector: torch.Tensor [d_model], token_ids, top_logits)
127
+ """
128
+ inputs = tokenizer(text, return_tensors="pt").to(device)
129
+ with torch.no_grad():
130
+ outputs = model(
131
+ **inputs,
132
+ output_hidden_states=True,
133
+ )
134
+ # hidden_states is tuple of (layer+1) x (batch, seq, d_model)
135
+ # index layer_idx gives the residual stream AFTER that layer
136
+ hidden = outputs.hidden_states[layer_idx] # [1, seq, d_model]
137
+ # take last token activation
138
+ act = hidden[0, -1, :] # [d_model]
139
+ # also get logits for top-k
140
+ logits = outputs.logits[0, -1, :] # [vocab]
141
+ top_vals, top_idxs = torch.topk(logits, k=10)
142
+ return act.cpu(), inputs["input_ids"][0], top_idxs.cpu(), top_vals.cpu()
143
+
144
+
145
+ def compute_activation_stats(activations: list):
146
+ """Compute L2 norm stats over a list of activation vectors."""
147
+ norms_list = [a.norm().item() for a in activations]
148
+ norms_t = torch.tensor(norms_list)
149
+ stats = {
150
+ "mean": norms_t.mean().item(),
151
+ "p50": norms_t.median().item(),
152
+ "p90": norms_t.kthvalue(max(1, int(len(norms_t) * 0.9))).values.item(),
153
+ "max": norms_t.max().item(),
154
+ "min": norms_t.min().item(),
155
+ "std": norms_t.std().item(),
156
+ }
157
+ return stats
158
+
159
+
160
+ # ── injection token search ─────────────────────────────
161
+ def find_injection_token(tokenizer) -> tuple:
162
+ """
163
+ Find a rare character that tokenizes as a single token.
164
+ Returns (char, token_id).
165
+ """
166
+ # Candidates: CJK rare chars that are likely single tokens
167
+ candidates = [
168
+ "㈎", "㈏", "㈐", "㈑", "㈒", "㈓", "㈔", "㈕", "㈖", "㈗",
169
+ "㈘", "㈙", "㈚", "㈛", "㈜", "㈝", "㈞",
170
+ "ⓐ", "ⓑ", "ⓒ", "ⓓ", "ⓔ", "ⓕ",
171
+ "🀀", "🀁", "🀂", "🀃", "🀄", "🀅",
172
+ "✪", "✫", "✬", "✭", "✮", "✯",
173
+ "〓", "■", "□", "◆", "◇",
174
+ ]
175
+
176
+ print("=" * 60)
177
+ print("🔤 Injection Token Search")
178
+ print("=" * 60)
179
+
180
+ for char in candidates:
181
+ encoded = tokenizer.encode(char, add_special_tokens=False)
182
+ if len(encoded) == 1:
183
+ tid = encoded[0]
184
+ # Verify roundtrip
185
+ decoded = tokenizer.decode([tid])
186
+ print(f" ✅ Found single token: {char!r} -> id={tid} (decode: {decoded!r})")
187
+ return char, tid
188
+
189
+ # Fallback: search for ANY single-token char more systematically
190
+ print(" ⚠️ Searching more broadly for single-token chars...")
191
+ for codepoint in range(0x4E00, 0x9FFF): # CJK Unified
192
+ char = chr(codepoint)
193
+ encoded = tokenizer.encode(char, add_special_tokens=False)
194
+ if len(encoded) == 1:
195
+ tid = encoded[0]
196
+ decoded = tokenizer.decode([tid])
197
+ print(f" ✅ Found: {char!r} (U+{codepoint:04X}) -> id={tid}")
198
+ return char, tid
199
+
200
+ raise RuntimeError("Could not find a single-token injection character!")
201
+
202
+
203
+ # ── injection smoke test ───────────────────────────────
204
+ def injection_smoke(
205
+ model, tokenizer, device, dtype,
206
+ layer_idx: int, injection_scale: float, injection_char: str, injection_token_id: int,
207
+ ):
208
+ """
209
+ Build a prompt, replace injection token embedding with scaled activation,
210
+ run forward + generation. Verify no crash and shape mismatch.
211
+ """
212
+ print("=" * 60)
213
+ print("🧪 Injection Smoke Test")
214
+ print("=" * 60)
215
+
216
+ prompt = f"Context: 这是一个测试句子。\n<concept>{injection_char}</concept>\n<explanation>"
217
+ print(f" Prompt: {prompt!r}")
218
+
219
+ # Tokenize
220
+ inputs = tokenizer(prompt, return_tensors="pt").to(device)
221
+ input_ids = inputs["input_ids"]
222
+ print(f" Input shape: {input_ids.shape}")
223
+
224
+ # Find the injection token position
225
+ inj_ids = (input_ids[0] == injection_token_id).nonzero(as_tuple=True)[0]
226
+ if len(inj_ids) == 0:
227
+ raise RuntimeError(f"Injection token {injection_char!r} (id={injection_token_id}) not found in prompt!")
228
+ inj_pos = inj_ids[0].item()
229
+ print(f" Injection token at position: {inj_pos}")
230
+
231
+ # Get a dummy activation vector: random vector scaled to injection_scale
232
+ model_config = model.config
233
+ d_model = model_config.hidden_size
234
+ dummy_act = torch.randn(d_model, device=device, dtype=dtype)
235
+ dummy_act = dummy_act / dummy_act.norm() * injection_scale
236
+
237
+ # Forward with input_embeds injection
238
+ with torch.no_grad():
239
+ # Get base embeddings
240
+ embeds = model.get_input_embeddings()(input_ids) # [1, seq, d_model]
241
+ # Replace the injection position's embedding with our activation
242
+ embeds[0, inj_pos, :] = dummy_act
243
+
244
+ # Forward with input_embeds
245
+ outputs = model(inputs_embeds=embeds, output_hidden_states=True)
246
+ logits = outputs.logits
247
+ print(f" Forward output shape: {logits.shape} ✅")
248
+
249
+ # Extract the hidden state at injection position from layer layer_idx
250
+ hidden = outputs.hidden_states[layer_idx]
251
+ injected_act = hidden[0, inj_pos, :]
252
+ print(f" Hidden at injection pos, layer {layer_idx}: shape={injected_act.shape} ✅")
253
+
254
+ # Compare input activation vs post-layer activation
255
+ cos_sim = F.cosine_similarity(dummy_act.unsqueeze(0), injected_act.unsqueeze(0))
256
+ print(f" Cosine(input_act, post_layer_{layer_idx}_act): {cos_sim.item():.4f}")
257
+
258
+ # Generation test — greedy
259
+ print("\n 🔄 Generation test (greedy, max 20 tokens)...")
260
+ gen_outputs = model.generate(
261
+ inputs_embeds=embeds,
262
+ max_new_tokens=20,
263
+ do_sample=False,
264
+ pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
265
+ )
266
+ gen_text = tokenizer.decode(gen_outputs[0], skip_special_tokens=True)
267
+ print(f" Generated: {gen_text!r}")
268
+ print(f" ✅ Injection generation completed without error")
269
+
270
+ return {
271
+ "forward_ok": True,
272
+ "gen_ok": True,
273
+ "cosine_post_injection": cos_sim.item(),
274
+ }
275
+
276
+
277
+ # ── main ───────────────────────────────────────────────
278
+ def main():
279
+ print("=" * 60)
280
+ print("🏗️ TINY-NLA STAGE 0 — Environment & Injection Smoke")
281
+ print("=" * 60)
282
+ print()
283
+
284
+ # 1. Device
285
+ device = detect_device()
286
+
287
+ # 2. MPS backward smoke
288
+ if not mps_backward_smoke(device):
289
+ device = torch.device("cpu")
290
+
291
+ # 3. dtype
292
+ dtype = get_dtype_for_device(device)
293
+ print(f" Using dtype: {dtype}")
294
+ print()
295
+
296
+ # 4. Load model
297
+ model, tokenizer = load_model_and_tokenizer(BASE_MODEL, device, dtype)
298
+
299
+ # 5. Inspect
300
+ info = inspect_model(model)
301
+ layer_idx = info["layer_index"]
302
+ d_model = info["hidden_size"]
303
+
304
+ # 6. Extract sample activations
305
+ print("=" * 60)
306
+ print("📊 Sample Activation Extraction")
307
+ print("=" * 60)
308
+ sample_texts = [
309
+ "人工智能正在改变世界。",
310
+ "深度学习模型可以学习复杂模式。",
311
+ "今天天气真好,适合出去散步。",
312
+ "Qwen3是一个强大的语言模型。",
313
+ ]
314
+ activations = []
315
+ for txt in sample_texts:
316
+ act, ids, topk_ids, topk_vals = extract_activation(model, tokenizer, txt, layer_idx, device)
317
+ activations.append(act)
318
+ top_tokens = [tokenizer.decode([t]) for t in topk_ids[:3]]
319
+ print(f" Text: {txt}")
320
+ print(f" Last token activation norm: {act.norm():.4f}")
321
+ print(f" Top-3 next tokens: {top_tokens}")
322
+
323
+ # 7. Activation stats
324
+ act_stats = compute_activation_stats(activations)
325
+ print()
326
+ print("📊 Activation L2 Norm Statistics")
327
+ print("-" * 40)
328
+ for k, v in act_stats.items():
329
+ print(f" {k}: {v:.4f}")
330
+
331
+ injection_scale = round(act_stats["p50"], 4)
332
+ print(f"\n → injection_scale = p50 = {injection_scale}")
333
+
334
+ # 8. Injection token search
335
+ inj_char, inj_token_id = find_injection_token(tokenizer)
336
+ print()
337
+
338
+ # 9. Injection smoke test
339
+ smoke_result = injection_smoke(
340
+ model, tokenizer, device, dtype,
341
+ layer_idx, injection_scale, inj_char, inj_token_id,
342
+ )
343
+ print()
344
+
345
+ # 10. Write nla_meta.yaml
346
+ meta = {
347
+ "kind": "tiny_nla_model",
348
+ "base_model": BASE_MODEL,
349
+ "av_init_model": INSTRUCT_MODEL,
350
+ "layer_index": layer_idx,
351
+ "num_hidden_layers": info["num_hidden_layers"],
352
+ "d_model": d_model,
353
+ "activation_source": "residual_stream",
354
+ "token_position_policy": "selected_token",
355
+ "extraction": {
356
+ "injection_scale": injection_scale,
357
+ "mse_normalization": "l2_direction",
358
+ },
359
+ "tokens": {
360
+ "injection_char": inj_char,
361
+ "injection_token_id": inj_token_id,
362
+ },
363
+ "prompt_templates": {
364
+ "av": f"<concept>{{injection_char}}</concept>\n<explanation>",
365
+ "ar": f"<explanation>{{explanation_text}}</explanation>",
366
+ },
367
+ "training": {
368
+ "device": device.type,
369
+ "dtype": str(dtype),
370
+ "dataset_size": 0,
371
+ "teacher": "local_qwen_instruct",
372
+ "created_at": time.strftime("%Y-%m-%dT%H:%M:%S"),
373
+ },
374
+ "activation_stats": act_stats,
375
+ "smoke_test": smoke_result,
376
+ }
377
+
378
+ with open(SIDECAR_PATH, "w", encoding="utf-8") as f:
379
+ yaml.dump(meta, f, allow_unicode=True, sort_keys=False, default_flow_style=False)
380
+ print(f" ✅ Sidecar written: {SIDECAR_PATH}")
381
+
382
+ # 11. Summary
383
+ print()
384
+ print("=" * 60)
385
+ print("✅ STAGE 0 COMPLETE")
386
+ print("=" * 60)
387
+ print(f" Device: {device.type}")
388
+ print(f" dtype: {dtype}")
389
+ print(f" Layers: {info['num_hidden_layers']}")
390
+ print(f" Selected layer: {layer_idx} ({layer_idx/info['num_hidden_layers']*100:.0f}% depth)")
391
+ print(f" d_model: {d_model}")
392
+ print(f" Injection char: {inj_char!r} (id={inj_token_id})")
393
+ print(f" Injection scale: {injection_scale}")
394
+ print(f" Forward smoke: {'✅ PASS' if smoke_result['forward_ok'] else '❌ FAIL'}")
395
+ print(f" Generation smoke: {'✅ PASS' if smoke_result['gen_ok'] else '❌ FAIL'}")
396
+ print(f" Sidecar: {SIDECAR_PATH}")
397
+ print()
398
+
399
+
400
+ if __name__ == "__main__":
401
+ main()
experiments/tiny_nla/train_ar.py ADDED
@@ -0,0 +1,301 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Stage 1: AR SFT — Activation Reconstructor
4
+
5
+ Input: explanation text (via teacher)
6
+ Output: reconstructed activation vector (d_model)
7
+
8
+ Architecture:
9
+ - Qwen3-0.6B-Base trunk (frozen)
10
+ - Linear(d_model, d_model) head on last token hidden state
11
+ - L2-normalized MSE loss
12
+
13
+ Baselines:
14
+ - Mean vector baseline (always predict mean activation)
15
+ - Shuffled label baseline (random pairing)
16
+ """
17
+
18
+ import json, os, time, math, yaml, random
19
+ from pathlib import Path
20
+
21
+ import torch
22
+ import torch.nn as nn
23
+ import torch.nn.functional as F
24
+ from torch.utils.data import Dataset, DataLoader
25
+ from transformers import AutoModelForCausalLM, AutoTokenizer
26
+
27
+ # ── paths ──────────────────────────────────────────────
28
+ REPO_ROOT = Path(__file__).resolve().parents[2]
29
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
30
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
31
+ CHECKPOINT_DIR = ARTIFACTS_DIR / "checkpoints" / "ar"
32
+ CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
33
+
34
+ BASE_MODEL = "Qwen/Qwen3-0.6B-Base"
35
+
36
+ # ── device ─────────────────────────────────────────────
37
+ def detect_device():
38
+ if torch.cuda.is_available():
39
+ return torch.device("cuda")
40
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
41
+ return torch.device("mps")
42
+ return torch.device("cpu")
43
+
44
+ # ── AR Head ────────────────────────────────────────────
45
+ class ARHead(nn.Module):
46
+ """Simple linear projection from d_model to d_model."""
47
+ def __init__(self, d_model: int):
48
+ super().__init__()
49
+ self.linear = nn.Linear(d_model, d_model, bias=False)
50
+
51
+ def forward(self, x):
52
+ # x: [batch, d_model] — last token hidden state
53
+ return self.linear(x)
54
+
55
+
56
+ # ── Dataset ────────────────────────────────────────────
57
+ class ARDataset(Dataset):
58
+ """Explanation text → activation vector."""
59
+ def __init__(self, records, activations, tokenizer, max_length=128):
60
+ self.tokenizer = tokenizer
61
+ self.max_length = max_length
62
+ self.data = []
63
+
64
+ for rec, act in zip(records, activations):
65
+ explanation = rec.get("teacher_explanation", "") or rec.get("teacher_explanation_raw", "")
66
+ if not explanation or explanation in ("[空输出]", ""):
67
+ continue
68
+ self.data.append((explanation, act))
69
+
70
+ def __len__(self):
71
+ return len(self.data)
72
+
73
+ def __getitem__(self, idx):
74
+ text, act = self.data[idx]
75
+ tokens = self.tokenizer(
76
+ text,
77
+ truncation=True,
78
+ padding="max_length",
79
+ max_length=self.max_length,
80
+ return_tensors="pt",
81
+ )
82
+ return {
83
+ "input_ids": tokens["input_ids"][0],
84
+ "attention_mask": tokens["attention_mask"][0],
85
+ "target": act, # [d_model]
86
+ }
87
+
88
+
89
+ # ── Loss ───────────────────────────────────────────────
90
+ def normalized_mse_loss(pred, target):
91
+ """
92
+ L2-normalize both, then MSE = 2*(1-cosine)
93
+ Equivalent to ||normalized_pred - normalized_target||²
94
+ """
95
+ pred_n = F.normalize(pred, dim=-1)
96
+ target_n = F.normalize(target, dim=-1)
97
+ return F.mse_loss(pred_n, target_n)
98
+
99
+
100
+ def cosine_similarity(pred, target):
101
+ pred_n = F.normalize(pred, dim=-1)
102
+ target_n = F.normalize(target, dim=-1)
103
+ return (pred_n * target_n).sum(dim=-1)
104
+
105
+
106
+ # ── Training ───────────────────────────────────────────
107
+ def train():
108
+ print("=" * 60)
109
+ print("🔧 AR SFT — Activation Reconstructor Training")
110
+ print("=" * 60)
111
+
112
+ device = detect_device()
113
+ dtype = torch.float32
114
+ print(f" Device: {device}, dtype: {dtype}")
115
+
116
+ # Load sidecar
117
+ with open(SIDECAR_PATH, "r") as f:
118
+ nla_meta = yaml.safe_load(f)
119
+ d_model = nla_meta["d_model"]
120
+ print(f" d_model: {d_model}")
121
+
122
+ # Load dataset
123
+ print("\n Loading dataset...")
124
+ jsonl_path = ARTIFACTS_DIR / "dataset.jsonl"
125
+ act_path = ARTIFACTS_DIR / "activations.pt"
126
+
127
+ records = []
128
+ with open(jsonl_path, "r", encoding="utf-8") as f:
129
+ for line in f:
130
+ records.append(json.loads(line))
131
+ activations = torch.load(act_path, weights_only=True)
132
+ print(f" Loaded {len(records)} records, activations shape: {activations.shape}")
133
+
134
+ # Train/val split
135
+ random.seed(42)
136
+ indices = list(range(len(records)))
137
+ random.shuffle(indices)
138
+ val_size = max(1, int(len(indices) * 0.15))
139
+ train_idx, val_idx = indices[val_size:], indices[:val_size]
140
+ print(f" Train: {len(train_idx)}, Val: {len(val_idx)}")
141
+
142
+ # Load tokenizer
143
+ tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
144
+ tokenizer.pad_token = tokenizer.eos_token
145
+
146
+ train_records = [records[i] for i in train_idx]
147
+ train_acts = activations[train_idx]
148
+ val_records = [records[i] for i in val_idx]
149
+ val_acts = activations[val_idx]
150
+
151
+ train_dataset = ARDataset(train_records, train_acts, tokenizer)
152
+ val_dataset = ARDataset(val_records, val_acts, tokenizer)
153
+ print(f" Train dataset: {len(train_dataset)}, Val dataset: {len(val_dataset)}")
154
+
155
+ train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
156
+ val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)
157
+
158
+ # Load base model (frozen)
159
+ print("\n Loading base model (frozen)...")
160
+ t0 = time.perf_counter()
161
+ base_model = AutoModelForCausalLM.from_pretrained(
162
+ BASE_MODEL,
163
+ trust_remote_code=True,
164
+ torch_dtype=dtype,
165
+ low_cpu_mem_usage=True,
166
+ attn_implementation="eager",
167
+ ).to(device)
168
+ base_model.eval()
169
+ for p in base_model.parameters():
170
+ p.requires_grad = False
171
+ print(f" Base model loaded in {time.perf_counter() - t0:.1f}s")
172
+
173
+ # AR head
174
+ ar_head = ARHead(d_model).to(device)
175
+ print(f" AR head params: {sum(p.numel() for p in ar_head.parameters()):,}")
176
+
177
+ optimizer = torch.optim.AdamW(ar_head.parameters(), lr=1e-3)
178
+
179
+ # ── Compute baselines ──
180
+ print("\n📊 Computing baselines...")
181
+
182
+ # Mean vector baseline
183
+ all_train_acts = torch.stack([train_dataset[i]["target"] for i in range(len(train_dataset))])
184
+ mean_vec = all_train_acts.mean(dim=0) # [d_model]
185
+
186
+ mean_cosines = []
187
+ for i in range(len(val_dataset)):
188
+ target = val_dataset[i]["target"]
189
+ cos = F.cosine_similarity(mean_vec.unsqueeze(0), target.unsqueeze(0))
190
+ mean_cosines.append(cos.item())
191
+ mean_baseline = torch.tensor(mean_cosines).mean().item()
192
+ print(f" Mean-vector baseline cosine: {mean_baseline:.4f}")
193
+
194
+ # Shuffled baseline
195
+ val_targets = torch.stack([val_dataset[i]["target"] for i in range(len(val_dataset))])
196
+ shuffled = val_targets[torch.randperm(len(val_targets))]
197
+ shuffled_cosines = F.cosine_similarity(val_targets, shuffled, dim=-1)
198
+ shuffled_baseline = shuffled_cosines.mean().item()
199
+ print(f" Shuffled baseline cosine: {shuffled_baseline:.4f}")
200
+
201
+ # ── Training loop ──
202
+ print("\n🏋️ Training AR head...")
203
+ n_epochs = 50
204
+ best_val_loss = float("inf")
205
+ best_val_cosine = 0.0
206
+
207
+ for epoch in range(n_epochs):
208
+ # Train
209
+ ar_head.train()
210
+ total_loss = 0.0
211
+ n_batches = 0
212
+
213
+ for batch in train_loader:
214
+ input_ids = batch["input_ids"].to(device)
215
+ attention_mask = batch["attention_mask"].to(device)
216
+ targets = batch["target"].to(device)
217
+
218
+ with torch.no_grad():
219
+ outputs = base_model(
220
+ input_ids=input_ids,
221
+ attention_mask=attention_mask,
222
+ output_hidden_states=True,
223
+ )
224
+ # Get last token hidden state from last layer
225
+ last_hidden = outputs.hidden_states[-1] # [batch, seq, d_model]
226
+ # Use the last non-pad token position
227
+ seq_lens = attention_mask.sum(dim=1) - 1
228
+ batch_indices = torch.arange(last_hidden.size(0), device=device)
229
+ last_token_hidden = last_hidden[batch_indices, seq_lens, :] # [batch, d_model]
230
+
231
+ pred = ar_head(last_token_hidden)
232
+ loss = normalized_mse_loss(pred, targets)
233
+
234
+ optimizer.zero_grad()
235
+ loss.backward()
236
+ optimizer.step()
237
+
238
+ total_loss += loss.item()
239
+ n_batches += 1
240
+
241
+ avg_train_loss = total_loss / n_batches
242
+
243
+ # Validation
244
+ ar_head.eval()
245
+ val_losses = []
246
+ val_cosines = []
247
+
248
+ with torch.no_grad():
249
+ for batch in val_loader:
250
+ input_ids = batch["input_ids"].to(device)
251
+ attention_mask = batch["attention_mask"].to(device)
252
+ targets = batch["target"].to(device)
253
+
254
+ outputs = base_model(
255
+ input_ids=input_ids,
256
+ attention_mask=attention_mask,
257
+ output_hidden_states=True,
258
+ )
259
+ last_hidden = outputs.hidden_states[-1]
260
+ seq_lens = attention_mask.sum(dim=1) - 1
261
+ batch_indices = torch.arange(last_hidden.size(0), device=device)
262
+ last_token_hidden = last_hidden[batch_indices, seq_lens, :]
263
+
264
+ pred = ar_head(last_token_hidden)
265
+ loss = normalized_mse_loss(pred, targets)
266
+ cos = cosine_similarity(pred, targets)
267
+
268
+ val_losses.append(loss.item())
269
+ val_cosines.extend(cos.cpu().tolist())
270
+
271
+ avg_val_loss = sum(val_losses) / len(val_losses)
272
+ avg_val_cosine = sum(val_cosines) / len(val_cosines)
273
+
274
+ if avg_val_cosine > best_val_cosine:
275
+ best_val_cosine = avg_val_cosine
276
+ best_val_loss = avg_val_loss
277
+ torch.save(ar_head.state_dict(), CHECKPOINT_DIR / "best_ar_head.pt")
278
+
279
+ if (epoch + 1) % 5 == 0 or epoch == 0:
280
+ print(f" Epoch {epoch+1:2d}/{n_epochs} | train_loss={avg_train_loss:.6f} | val_loss={avg_val_loss:.6f} | val_cos={avg_val_cosine:.4f} | best_cos={best_val_cosine:.4f}")
281
+
282
+ # ── Final evaluation ──
283
+ print(f"\n📈 Final Results")
284
+ print(f" {'':30} {'Loss':>10} {'Cosine':>10}")
285
+ print(f" {'Mean-vector baseline':30} {'':>10} {mean_baseline:>10.4f}")
286
+ print(f" {'Shuffled baseline':30} {'':>10} {shuffled_baseline:>10.4f}")
287
+ print(f" {'AR trained (best)':30} {best_val_loss:>10.6f} {best_val_cosine:>10.4f}")
288
+
289
+ improvement = best_val_cosine - max(mean_baseline, shuffled_baseline)
290
+ print(f"\n Improvement over best baseline: {improvement:+.4f}")
291
+
292
+ if improvement <= 0:
293
+ print(" ⚠️ AR does NOT beat baselines! May need fixes.")
294
+
295
+ print(f"\n Checkpoint: {CHECKPOINT_DIR / 'best_ar_head.pt'}")
296
+
297
+ return best_val_cosine, best_val_loss
298
+
299
+
300
+ if __name__ == "__main__":
301
+ train()
experiments/tiny_nla/train_av.py ADDED
@@ -0,0 +1,331 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Stage 2: AV SFT — Activation Verbalizer
4
+
5
+ Fixes vs prior version:
6
+ 1. No padding='max_length' — pad to batch max length with attention_mask
7
+ 2. Pass attention_mask in both train forward and generate
8
+ 3. Generate with prompt-only embeds (not full padded sequence)
9
+ 4. Use attention_mask in AVModel.forward
10
+ """
11
+
12
+ import json, os, time, yaml, random
13
+ from pathlib import Path
14
+
15
+ import torch
16
+ import torch.nn as nn
17
+ from torch.utils.data import Dataset, DataLoader
18
+ from transformers import AutoModelForCausalLM, AutoTokenizer
19
+ from peft import LoraConfig, get_peft_model, TaskType
20
+
21
+ REPO_ROOT = Path(__file__).resolve().parents[2]
22
+ SIDECAR_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
23
+ ARTIFACTS_DIR = REPO_ROOT / "artifacts" / "tiny_nla"
24
+ CHECKPOINT_DIR = ARTIFACTS_DIR / "checkpoints" / "av"
25
+ CHECKPOINT_DIR.mkdir(parents=True, exist_ok=True)
26
+
27
+ INSTRUCT_MODEL = "Qwen/Qwen3-0.6B"
28
+
29
+
30
+ def detect_device():
31
+ if torch.cuda.is_available():
32
+ return torch.device("cuda")
33
+ if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
34
+ return torch.device("mps")
35
+ return torch.device("cpu")
36
+
37
+
38
+ def collate_fn(batch, pad_id):
39
+ """Pad to batch max length (not global max)."""
40
+ max_len = max(b["input_ids"].shape[0] for b in batch)
41
+ input_ids = torch.full((len(batch), max_len), pad_id, dtype=torch.long)
42
+ attention_mask = torch.zeros(len(batch), max_len, dtype=torch.long)
43
+ labels = torch.full((len(batch), max_len), -100, dtype=torch.long)
44
+ activations = torch.stack([b["activation"] for b in batch])
45
+ prompt_lens = [b["prompt_len"] for b in batch]
46
+
47
+ for i, b in enumerate(batch):
48
+ n = b["input_ids"].shape[0]
49
+ input_ids[i, :n] = b["input_ids"]
50
+ attention_mask[i, :n] = 1
51
+ labels[i, :n] = b["labels"]
52
+
53
+ return {
54
+ "input_ids": input_ids,
55
+ "attention_mask": attention_mask,
56
+ "labels": labels,
57
+ "activations": activations,
58
+ "prompt_lens": prompt_lens,
59
+ }
60
+
61
+
62
+ class AVDataset(Dataset):
63
+ def __init__(self, av_data, activations, tokenizer, inj_char, inj_token_id, inj_scale, max_length=96):
64
+ self.tokenizer = tokenizer
65
+ self.inj_token_id = inj_token_id
66
+ self.inj_scale = inj_scale
67
+ self.max_length = max_length
68
+ self.inj_char = inj_char
69
+ self.data = []
70
+
71
+ for rec, act in zip(av_data, activations):
72
+ expl = rec.get("teacher_explanation", "") or ""
73
+ if not expl.strip() or expl in ("[空输出]",):
74
+ continue
75
+ self.data.append((expl, act))
76
+
77
+ print(f" AVDataset: {len(self.data)} samples")
78
+
79
+ def __len__(self):
80
+ return len(self.data)
81
+
82
+ def __getitem__(self, idx):
83
+ explanation, activation = self.data[idx]
84
+ prompt = f"<concept>{self.inj_char}</concept>\n<explanation>"
85
+
86
+ # Tokenize prompt and full sequence separately (no padding here)
87
+ prompt_ids = self.tokenizer(prompt, return_tensors="pt")["input_ids"][0]
88
+ expl_ids = self.tokenizer(
89
+ explanation,
90
+ add_special_tokens=False,
91
+ return_tensors="pt",
92
+ )["input_ids"][0]
93
+
94
+ # Truncate explanation if needed
95
+ max_expl = self.max_length - len(prompt_ids) - 1 # -1 for eos
96
+ expl_ids = expl_ids[:max_expl]
97
+
98
+ # Build full sequence: prompt + explanation + eos
99
+ eos = torch.tensor([self.tokenizer.eos_token_id], dtype=torch.long)
100
+ input_ids = torch.cat([prompt_ids, expl_ids, eos])
101
+
102
+ # Labels: -100 for prompt, actual ids for explanation+eos
103
+ labels = torch.full_like(input_ids, -100)
104
+ labels[len(prompt_ids):] = torch.cat([expl_ids, eos])
105
+
106
+ return {
107
+ "input_ids": input_ids,
108
+ "labels": labels,
109
+ "activation": activation,
110
+ "prompt_len": len(prompt_ids),
111
+ }
112
+
113
+
114
+ class AVModel(nn.Module):
115
+ def __init__(self, base_model, inj_token_id):
116
+ super().__init__()
117
+ self.model = base_model
118
+ self.inj_token_id = inj_token_id
119
+
120
+ def forward(self, input_ids, attention_mask, labels, activations):
121
+ embeds = self.model.get_input_embeddings()(input_ids)
122
+
123
+ for b in range(input_ids.shape[0]):
124
+ positions = (input_ids[b] == self.inj_token_id).nonzero(as_tuple=True)[0]
125
+ if len(positions) > 0:
126
+ embeds[b, positions[0].item(), :] = activations[b].to(embeds.dtype)
127
+
128
+ return self.model(
129
+ inputs_embeds=embeds,
130
+ attention_mask=attention_mask,
131
+ labels=labels,
132
+ )
133
+
134
+
135
+ def scale_activations(acts, scale, device):
136
+ acts = acts.to(device)
137
+ norms = acts.norm(dim=-1, keepdim=True).clamp(min=1e-6)
138
+ return acts / norms * scale
139
+
140
+
141
+ def train():
142
+ print("=" * 60)
143
+ print("AV SFT — Activation Verbalizer Training")
144
+ print("=" * 60)
145
+
146
+ device = detect_device()
147
+ dtype = torch.float32
148
+ print(f" Device: {device}, dtype: {dtype}")
149
+
150
+ with open(SIDECAR_PATH) as f:
151
+ meta = yaml.safe_load(f)
152
+
153
+ inj_char = meta["tokens"]["injection_char"]
154
+ inj_token_id = meta["tokens"]["injection_token_id"]
155
+ inj_scale = meta["extraction"]["injection_scale"]
156
+ print(f" inj_char={inj_char!r} id={inj_token_id} scale={inj_scale}")
157
+
158
+ with open(ARTIFACTS_DIR / "av_training_data.json", encoding="utf-8") as f:
159
+ av_data = json.load(f)
160
+ activations = torch.load(ARTIFACTS_DIR / "av_activations.pt", weights_only=True)
161
+ print(f" Loaded {len(av_data)} records, activations {activations.shape}")
162
+
163
+ tokenizer = AutoTokenizer.from_pretrained(INSTRUCT_MODEL, trust_remote_code=True)
164
+ # Use a different pad token to avoid pad=eos issue
165
+ tokenizer.pad_token_id = tokenizer.eos_token_id # needed for generate only
166
+
167
+ random.seed(42)
168
+ idx = list(range(len(av_data)))
169
+ random.shuffle(idx)
170
+ val_n = max(20, int(len(idx) * 0.15))
171
+ train_idx, val_idx = idx[val_n:], idx[:val_n]
172
+
173
+ def mk_dataset(indices):
174
+ return AVDataset(
175
+ [av_data[i] for i in indices],
176
+ activations[indices],
177
+ tokenizer, inj_char, inj_token_id, inj_scale,
178
+ )
179
+
180
+ train_ds = mk_dataset(train_idx)
181
+ val_ds = mk_dataset(val_idx)
182
+
183
+ pad_id = tokenizer.eos_token_id
184
+ train_loader = DataLoader(train_ds, batch_size=4, shuffle=True,
185
+ collate_fn=lambda b: collate_fn(b, pad_id))
186
+ val_loader = DataLoader(val_ds, batch_size=4, shuffle=False,
187
+ collate_fn=lambda b: collate_fn(b, pad_id))
188
+
189
+ print(f"\n Loading {INSTRUCT_MODEL}...")
190
+ base_model = AutoModelForCausalLM.from_pretrained(
191
+ INSTRUCT_MODEL, trust_remote_code=True, dtype=dtype,
192
+ low_cpu_mem_usage=True, attn_implementation="eager",
193
+ ).to(device)
194
+
195
+ lora_cfg = LoraConfig(
196
+ task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=16,
197
+ lora_dropout=0.1, target_modules=["q_proj", "v_proj"],
198
+ bias="none",
199
+ )
200
+ lora_model = get_peft_model(base_model, lora_cfg)
201
+ lora_model.print_trainable_parameters()
202
+
203
+ av_model = AVModel(lora_model, inj_token_id).to(device)
204
+ optimizer = torch.optim.AdamW(av_model.parameters(), lr=1e-4, weight_decay=0.05)
205
+ scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=50, eta_min=5e-6)
206
+
207
+ n_epochs = 50
208
+ best_val_loss = float("inf")
209
+ patience, patience_count = 10, 0
210
+
211
+ for epoch in range(n_epochs):
212
+ av_model.train()
213
+ train_losses = []
214
+ for batch in train_loader:
215
+ input_ids = batch["input_ids"].to(device)
216
+ attention_mask = batch["attention_mask"].to(device)
217
+ labels = batch["labels"].to(device)
218
+ scaled_acts = scale_activations(batch["activations"].float(), inj_scale, device)
219
+
220
+ out = av_model(input_ids, attention_mask, labels, scaled_acts)
221
+ loss = out.loss
222
+
223
+ optimizer.zero_grad()
224
+ loss.backward()
225
+ torch.nn.utils.clip_grad_norm_(av_model.parameters(), 1.0)
226
+ optimizer.step()
227
+ train_losses.append(loss.item())
228
+
229
+ av_model.eval()
230
+ val_losses = []
231
+ with torch.no_grad():
232
+ for batch in val_loader:
233
+ out = av_model(
234
+ batch["input_ids"].to(device),
235
+ batch["attention_mask"].to(device),
236
+ batch["labels"].to(device),
237
+ scale_activations(batch["activations"].float(), inj_scale, device),
238
+ )
239
+ val_losses.append(out.loss.item())
240
+
241
+ tl = sum(train_losses) / len(train_losses)
242
+ vl = sum(val_losses) / len(val_losses)
243
+
244
+ if vl < best_val_loss:
245
+ best_val_loss = vl
246
+ patience_count = 0
247
+ av_model.model.save_pretrained(CHECKPOINT_DIR)
248
+ tokenizer.save_pretrained(CHECKPOINT_DIR)
249
+ else:
250
+ patience_count += 1
251
+
252
+ scheduler.step()
253
+
254
+ if (epoch + 1) % 5 == 0 or epoch == 0:
255
+ print(f" Epoch {epoch+1:2d}/{n_epochs} | train={tl:.4f} | val={vl:.4f} | best={best_val_loss:.4f}")
256
+
257
+ if patience_count >= patience:
258
+ print(f" Early stop at epoch {epoch+1}")
259
+ break
260
+
261
+ print(f"\n Best val loss: {best_val_loss:.4f}")
262
+
263
+ # ── Generate examples (prompt-only embeds, no padding) ──
264
+ print("\n Generating held-out examples...")
265
+ av_model.eval()
266
+ examples = []
267
+
268
+ # Reload best checkpoint
269
+ from peft import PeftModel
270
+ best_base = AutoModelForCausalLM.from_pretrained(
271
+ INSTRUCT_MODEL, trust_remote_code=True, dtype=dtype,
272
+ low_cpu_mem_usage=True, attn_implementation="eager",
273
+ ).to(device)
274
+ best_model = PeftModel.from_pretrained(best_base, CHECKPOINT_DIR).to(device)
275
+ best_model.eval()
276
+
277
+ prompt_template = f"<concept>{inj_char}</concept>\n<explanation>"
278
+
279
+ with torch.no_grad():
280
+ for i in range(min(25, len(val_ds))):
281
+ item = val_ds[i]
282
+ act = item["activation"].unsqueeze(0).float()
283
+ scaled = scale_activations(act, inj_scale, device)
284
+
285
+ # Tokenize prompt only (no padding)
286
+ p_ids = tokenizer(prompt_template, return_tensors="pt")["input_ids"].to(device)
287
+ p_mask = torch.ones_like(p_ids)
288
+
289
+ embeds = best_model.get_input_embeddings()(p_ids)
290
+ inj_pos = (p_ids[0] == inj_token_id).nonzero(as_tuple=True)[0]
291
+ if len(inj_pos) > 0:
292
+ embeds[0, inj_pos[0].item(), :] = scaled[0].to(embeds.dtype)
293
+
294
+ out_ids = best_model.generate(
295
+ inputs_embeds=embeds,
296
+ attention_mask=p_mask,
297
+ max_new_tokens=80,
298
+ do_sample=False,
299
+ eos_token_id=tokenizer.eos_token_id,
300
+ pad_token_id=tokenizer.eos_token_id,
301
+ )
302
+ # output_ids when using inputs_embeds: only generated tokens
303
+ gen_text = tokenizer.decode(out_ids[0], skip_special_tokens=True).strip()
304
+
305
+ teacher = tokenizer.decode(
306
+ item["labels"][item["prompt_len"]:][item["labels"][item["prompt_len"]:] != -100],
307
+ skip_special_tokens=True,
308
+ )
309
+ examples.append({
310
+ "index": i,
311
+ "teacher": teacher,
312
+ "av_generated": gen_text,
313
+ })
314
+
315
+ nonempty = sum(1 for e in examples if e["av_generated"].strip())
316
+ print(f" Non-empty: {nonempty}/{len(examples)}")
317
+ for e in examples[:5]:
318
+ print(f" T: {e['teacher'][:60]}")
319
+ print(f" G: {e['av_generated'][:80]}")
320
+ print()
321
+
322
+ out_path = ARTIFACTS_DIR / "av_examples.json"
323
+ with open(out_path, "w", encoding="utf-8") as f:
324
+ json.dump(examples, f, ensure_ascii=False, indent=2)
325
+ print(f" Saved to {out_path}")
326
+ print(f"\n LoRA adapter: {CHECKPOINT_DIR}")
327
+ return best_val_loss
328
+
329
+
330
+ if __name__ == "__main__":
331
+ train()
experiments/tiny_nla/train_rl_grpo.py ADDED
@@ -0,0 +1,404 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Phase 3: GRPO RL — AV+AR joint training.
4
+ Reward = -MSE(L2_norm(original), L2_norm(reconstructed))
5
+ Monitor: tensorboard --logdir artifacts/tiny_nla/runs
6
+
7
+ Usage:
8
+ python train_rl_grpo.py [--steps 1000] [--group-size 4] [--batch 8]
9
+ """
10
+ import json, yaml, random, time, argparse, os
11
+ from pathlib import Path
12
+ from datetime import datetime
13
+
14
+ # Disable HuggingFace network access — use local cache only, no HEAD checks to huggingface.co
15
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
16
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
17
+ os.environ.setdefault("HF_DATASETS_OFFLINE", "1")
18
+
19
+ import torch
20
+ import torch.nn as nn
21
+ import torch.nn.functional as F
22
+ import pyarrow.parquet as pq
23
+ from torch.utils.tensorboard import SummaryWriter
24
+ from transformers import AutoModelForCausalLM, AutoTokenizer
25
+ from peft import PeftModel, LoraConfig, get_peft_model, TaskType
26
+
27
+ REPO_ROOT = Path(__file__).resolve().parents[2]
28
+ ARTIFACTS = REPO_ROOT / "artifacts" / "tiny_nla"
29
+ META_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
30
+ RUNS_DIR = ARTIFACTS / "runs"
31
+
32
+ meta = yaml.safe_load(open(META_PATH))
33
+ D_MODEL = meta["d_model"]
34
+ LAYER = meta["layer_index"]
35
+ INJ_CHAR = meta["tokens"]["injection_char"]
36
+ INJ_TOK_ID = meta["tokens"]["injection_token_id"]
37
+ INJ_SCALE = meta["extraction"]["injection_scale"]
38
+ BASE_MODEL = meta["base_model"]
39
+ INST_MODEL = meta["av_init_model"]
40
+
41
+
42
+ def get_device():
43
+ if torch.cuda.is_available(): return torch.device("cuda")
44
+ if torch.backends.mps.is_available(): return torch.device("mps")
45
+ return torch.device("cpu")
46
+
47
+
48
+ def log(msg):
49
+ print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)
50
+
51
+
52
+ def scale_acts(acts, dev):
53
+ a = acts.to(dev).float()
54
+ return a / a.norm(dim=-1, keepdim=True).clamp(1e-6) * INJ_SCALE
55
+
56
+
57
+ def norm_acts(acts, dev):
58
+ a = acts.to(dev).float()
59
+ return a / a.norm(dim=-1, keepdim=True).clamp(1e-6)
60
+
61
+
62
+ # ── Load activations only (no teacher needed for RL) ─
63
+ def load_activations():
64
+ table = pq.read_table(ARTIFACTS / "activations_v2.parquet")
65
+ acts = torch.tensor([table["activation"][i].as_py()
66
+ for i in range(len(table))], dtype=torch.float32)
67
+ log(f"Loaded {len(acts)} activations | norm {acts.norm(dim=-1).min():.1f}–{acts.norm(dim=-1).max():.1f}")
68
+ return acts
69
+
70
+
71
+ # ── AR head (frozen during RL, used as reward scorer) ─
72
+ class ARHead(nn.Module):
73
+ def __init__(self, d_model):
74
+ super().__init__()
75
+ self.proj = nn.Linear(d_model, d_model, bias=False)
76
+ def forward(self, h):
77
+ return F.normalize(self.proj(h), dim=-1)
78
+
79
+
80
+ def load_ar(dev):
81
+ ckpt = torch.load(ARTIFACTS / "checkpoints" / "ar_v2" / "ar_head_v2.pt",
82
+ map_location=dev, weights_only=True)
83
+ head = ARHead(D_MODEL).to(dev)
84
+ head.load_state_dict(ckpt["head"])
85
+ head.eval()
86
+ for p in head.parameters(): p.requires_grad_(False)
87
+ log(f"AR head loaded (val_cos={ckpt.get('val_cosine',0):.4f})")
88
+ return head
89
+
90
+
91
+ # ── AV: generate explanations for a batch of activations ─
92
+ def av_generate(av_model, tok, acts_scaled, dev, max_new=64, temperature=1.0, num_return=1):
93
+ """
94
+ acts_scaled: [B, D] already scaled to INJ_SCALE
95
+ Returns list of B lists, each with num_return explanation strings.
96
+ """
97
+ prompt = f"<concept>{INJ_CHAR}</concept>\n<explanation>"
98
+ p_ids = tok(prompt, return_tensors="pt")["input_ids"].to(dev)
99
+ p_mask = torch.ones(1, p_ids.shape[1], device=dev, dtype=torch.long)
100
+ inj_pos = (p_ids[0] == INJ_TOK_ID).nonzero(as_tuple=True)[0][0].item()
101
+
102
+ all_expls = []
103
+ for b in range(acts_scaled.shape[0]):
104
+ embeds = av_model.get_input_embeddings()(p_ids).clone() # [1, P, D]
105
+ embeds[0, inj_pos] = acts_scaled[b].to(embeds.dtype)
106
+ expls = []
107
+ for _ in range(num_return):
108
+ with torch.no_grad():
109
+ out = av_model.generate(
110
+ inputs_embeds=embeds, attention_mask=p_mask,
111
+ max_new_tokens=max_new,
112
+ do_sample=(temperature > 0),
113
+ temperature=temperature if temperature > 0 else 1.0,
114
+ pad_token_id=tok.eos_token_id,
115
+ )
116
+ expls.append(tok.decode(out[0], skip_special_tokens=True).strip())
117
+ all_expls.append(expls)
118
+ return all_expls
119
+
120
+
121
+ # ── AR reconstruct explanations → activation ─────────
122
+ def ar_reconstruct(ar_backbone, ar_head, tok, explanations, dev):
123
+ """
124
+ explanations: list of strings
125
+ Returns: [N, D] normalized reconstructions (float32 for stable cosine)
126
+ """
127
+ tok.pad_token_id = tok.eos_token_id
128
+ enc = tok(explanations, return_tensors="pt", padding=True,
129
+ truncation=True, max_length=128).to(dev)
130
+ with torch.no_grad():
131
+ h = ar_backbone(**enc, output_hidden_states=True).hidden_states[-1]
132
+ lens = enc["attention_mask"].sum(1) - 1
133
+ last = h[torch.arange(len(h)), lens]
134
+ recon = ar_head(last)
135
+ return recon.float() # cast to float32 for downstream cosine
136
+
137
+
138
+ # ── GRPO loss ─────────────────────────────────────────
139
+ def grpo_loss(av_model, tok, acts_scaled, acts_norm, ar_backbone, ar_head,
140
+ dev, group_size, max_new, kl_coef, ref_av, temperature=1.2):
141
+ """
142
+ Returns: (policy_loss, mean_reward, reward_std)
143
+ """
144
+ B = acts_scaled.shape[0]
145
+ prompt = f"<concept>{INJ_CHAR}</concept>\n<explanation>"
146
+ p_ids = tok(prompt, return_tensors="pt")["input_ids"].to(dev)
147
+ inj_pos = (p_ids[0] == INJ_TOK_ID).nonzero(as_tuple=True)[0][0].item()
148
+
149
+ all_log_probs = [] # [B, G]
150
+ all_rewards = [] # [B, G]
151
+ ref_log_probs = [] # [B, G] for KL
152
+
153
+ for b in range(B):
154
+ embeds_base = av_model.get_input_embeddings()(p_ids).clone()
155
+ embeds_base[0, inj_pos] = acts_scaled[b].to(embeds_base.dtype)
156
+
157
+ b_lp, b_rew, b_rlp = [], [], []
158
+ for g in range(group_size):
159
+ # Sample with temperature
160
+ with torch.no_grad():
161
+ out = av_model.generate(
162
+ inputs_embeds=embeds_base.clone(),
163
+ attention_mask=torch.ones(1, p_ids.shape[1], device=dev),
164
+ max_new_tokens=max_new, do_sample=True, temperature=temperature,
165
+ pad_token_id=tok.eos_token_id, return_dict_in_generate=True,
166
+ output_scores=True,
167
+ )
168
+ gen_ids = out.sequences[0] # [T]
169
+ expl = tok.decode(gen_ids, skip_special_tokens=True).strip()
170
+
171
+ # Reward: AR roundtrip cosine
172
+ recon = ar_reconstruct(ar_backbone, ar_head, tok, [expl], dev) # [1,D]
173
+ orig_n = acts_norm[b].to(dev).unsqueeze(0)
174
+ reward = (recon * orig_n).sum(-1).item() # cosine similarity
175
+ b_rew.append(reward)
176
+
177
+ # Policy log-prob of generated tokens (recompute with grad)
178
+ full_embeds = av_model.get_input_embeddings()(p_ids).clone()
179
+ full_embeds[0, inj_pos] = acts_scaled[b].to(full_embeds.dtype)
180
+ # Append gen_ids to prompt via input_ids for teacher forcing
181
+ # (gen_ids are new tokens only from generate with inputs_embeds)
182
+ if len(gen_ids) > 0:
183
+ gen_embeds = av_model.get_input_embeddings()(gen_ids.unsqueeze(0))
184
+ all_embeds = torch.cat([full_embeds, gen_embeds], dim=1)
185
+ attn = torch.ones(1, all_embeds.shape[1], device=dev)
186
+ logits = av_model(inputs_embeds=all_embeds, attention_mask=attn).logits
187
+ # log-prob of generated part
188
+ P = p_ids.shape[1]
189
+ lp = 0.0
190
+ for t, tok_id in enumerate(gen_ids):
191
+ lp = lp + F.log_softmax(logits[0, P+t-1], dim=-1)[tok_id]
192
+ lp = lp / max(len(gen_ids), 1)
193
+ else:
194
+ lp = torch.tensor(0.0, device=dev)
195
+ b_lp.append(lp)
196
+
197
+ # Reference log-prob (KL penalty)
198
+ with torch.no_grad():
199
+ ref_logits = ref_av(inputs_embeds=all_embeds.detach(), attention_mask=attn).logits
200
+ rlp = sum(F.log_softmax(ref_logits[0, P+t-1], dim=-1)[tok_id]
201
+ for t, tok_id in enumerate(gen_ids))
202
+ rlp = rlp / max(len(gen_ids), 1)
203
+ b_rlp.append(rlp if isinstance(rlp, float) else rlp.item())
204
+
205
+ all_rewards.append(b_rew)
206
+ all_log_probs.append(b_lp)
207
+ ref_log_probs.append(b_rlp)
208
+
209
+ # GRPO: normalize rewards within group
210
+ loss = torch.tensor(0.0, device=dev, requires_grad=True)
211
+ flat_rewards = [r for group in all_rewards for r in group]
212
+ mean_rew = sum(flat_rewards) / len(flat_rewards)
213
+
214
+ for b in range(B):
215
+ rewards = all_rewards[b]
216
+ mu = sum(rewards) / group_size
217
+ std = (sum((r-mu)**2 for r in rewards)/group_size)**0.5 + 1e-6
218
+ for g in range(group_size):
219
+ adv = (rewards[g] - mu) / std
220
+ kl = all_log_probs[b][g] - ref_log_probs[b][g]
221
+ loss = loss - (adv * all_log_probs[b][g] - kl_coef * kl)
222
+
223
+ return loss / (B * group_size), mean_rew, (sum((r-mean_rew)**2 for r in flat_rewards)/len(flat_rewards))**0.5
224
+
225
+
226
+ def main():
227
+ parser = argparse.ArgumentParser()
228
+ parser.add_argument("--steps", type=int, default=1000)
229
+ parser.add_argument("--group-size", type=int, default=8)
230
+ parser.add_argument("--batch", type=int, default=4)
231
+ parser.add_argument("--max-new", type=int, default=64)
232
+ parser.add_argument("--kl-coef", type=float, default=0.05)
233
+ parser.add_argument("--lr-av", type=float, default=2e-5)
234
+ parser.add_argument("--temperature", type=float, default=1.2,
235
+ help="Sampling temperature for exploration (higher = more diverse)")
236
+ parser.add_argument("--eval-every", type=int, default=50)
237
+ parser.add_argument("--save-every", type=int, default=100)
238
+ parser.add_argument("--resume", action="store_true",
239
+ help="Resume RL from latest checkpoint")
240
+ args = parser.parse_args()
241
+
242
+ dev = get_device()
243
+ run_name = f"rl_{datetime.now().strftime('%m%d_%H%M')}"
244
+ writer = SummaryWriter(RUNS_DIR / run_name)
245
+ RUNS_DIR.mkdir(parents=True, exist_ok=True)
246
+
247
+ log(f"GRPO RL | device={dev} | steps={args.steps} | G={args.group_size} | B={args.batch} | T={args.temperature} | kl={args.kl_coef} | lr={args.lr_av}")
248
+ log(f"Monitor: tensorboard --logdir {RUNS_DIR}")
249
+
250
+ acts = load_activations()
251
+ acts_n = norm_acts(acts, "cpu")
252
+
253
+ # Load AR backbone (frozen scorer)
254
+ log(f"Loading AR backbone ({BASE_MODEL}, float16, sdpa)...")
255
+ ar_backbone = AutoModelForCausalLM.from_pretrained(
256
+ BASE_MODEL, trust_remote_code=True, dtype=torch.float16,
257
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(dev)
258
+ ar_backbone.eval()
259
+ for p in ar_backbone.parameters(): p.requires_grad_(False)
260
+ ar_head = load_ar(dev)
261
+ ar_head = ar_head.half() # match backbone dtype
262
+ tok_ar = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
263
+ tok_ar.pad_token_id = tok_ar.eos_token_id
264
+
265
+ # Load AV from SFT or resume checkpoint
266
+ rl_resume_dir = ARTIFACTS / "checkpoints" / "av_rl_resume"
267
+ if args.resume and (rl_resume_dir / "training_state.pt").exists():
268
+ av_ckpt = rl_resume_dir
269
+ log(f"Resuming AV from {av_ckpt}...")
270
+ else:
271
+ if args.resume:
272
+ log(f"⚠️ No RL resume checkpoint found, starting from SFT")
273
+ av_ckpt = ARTIFACTS / "checkpoints" / "av_v2"
274
+ log(f"Loading AV from {av_ckpt}...")
275
+ av_base = AutoModelForCausalLM.from_pretrained(
276
+ INST_MODEL, trust_remote_code=True, dtype=torch.float16,
277
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(dev)
278
+ av_model = PeftModel.from_pretrained(av_base, av_ckpt, is_trainable=True).to(dev)
279
+ tok_av = AutoTokenizer.from_pretrained(av_ckpt, trust_remote_code=True)
280
+ tok_av.pad_token_id = tok_av.eos_token_id
281
+
282
+ # Reference model (frozen copy for KL)
283
+ ref_base = AutoModelForCausalLM.from_pretrained(
284
+ INST_MODEL, trust_remote_code=True, dtype=torch.float16,
285
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(dev)
286
+ ref_av = PeftModel.from_pretrained(ref_base, av_ckpt).to(dev)
287
+ ref_av.eval()
288
+ for p in ref_av.parameters(): p.requires_grad_(False)
289
+
290
+ opt = torch.optim.AdamW(av_model.parameters(), lr=args.lr_av)
291
+
292
+ # SFT baseline or resume state
293
+ start_step = 1
294
+ if args.resume and (rl_resume_dir / "training_state.pt").exists():
295
+ ts = torch.load(rl_resume_dir / "training_state.pt",
296
+ map_location=dev, weights_only=False)
297
+ opt.load_state_dict(ts["optimizer"])
298
+ start_step = ts["step"] + 1
299
+ best_cos = ts["best_cos"]
300
+ best_step = ts["best_step"]
301
+ sft_cos = ts["sft_cos"]
302
+ log(f"Resumed from step {start_step} | best_cos={best_cos:.4f} at step {best_step} | sft_cos={sft_cos:.4f}")
303
+ writer.add_scalar("rl/sft_baseline", sft_cos, 0)
304
+ else:
305
+ log("Computing SFT baseline roundtrip...")
306
+ sft_cosines = []
307
+ with torch.no_grad():
308
+ sample_idx = random.sample(range(len(acts)), min(50, len(acts)))
309
+ for i in sample_idx:
310
+ act_s = scale_acts(acts[i:i+1], dev)
311
+ act_n = acts_n[i:i+1].to(dev)
312
+ expls = av_generate(av_model, tok_av, act_s, dev, args.max_new, temperature=0, num_return=1)
313
+ recon = ar_reconstruct(ar_backbone, ar_head, tok_ar, expls[0], dev)
314
+ sft_cosines.append((recon * act_n).sum(-1).item())
315
+ sft_cos = sum(sft_cosines)/len(sft_cosines)
316
+ log(f"SFT baseline roundtrip cosine: {sft_cos:.4f}")
317
+ writer.add_scalar("rl/sft_baseline", sft_cos, 0)
318
+ best_cos, best_step = sft_cos, 0
319
+
320
+ random.seed(42)
321
+
322
+ for step in range(start_step, args.steps+1):
323
+ av_model.train()
324
+ batch_idx = random.sample(range(len(acts)), args.batch)
325
+ batch_acts = acts[batch_idx]
326
+ batch_acts_n = acts_n[batch_idx]
327
+ acts_scaled = scale_acts(batch_acts, dev)
328
+
329
+ loss, mean_rew, rew_std = grpo_loss(
330
+ av_model, tok_av, acts_scaled, batch_acts_n,
331
+ ar_backbone, ar_head, dev,
332
+ args.group_size, args.max_new, args.kl_coef, ref_av,
333
+ temperature=args.temperature,
334
+ )
335
+
336
+ opt.zero_grad()
337
+ loss.backward()
338
+ torch.nn.utils.clip_grad_norm_(av_model.parameters(), 1.0)
339
+ opt.step()
340
+
341
+ writer.add_scalar("rl/loss", loss.item(), step)
342
+ writer.add_scalar("rl/mean_reward", mean_rew, step)
343
+ writer.add_scalar("rl/reward_std", rew_std, step)
344
+
345
+ if step % 10 == 0:
346
+ log(f"Step {step:4d}/{args.steps} | loss={loss.item():.4f} | "
347
+ f"reward={mean_rew:.4f}±{rew_std:.4f}")
348
+
349
+ # Eval roundtrip
350
+ if step % args.eval_every == 0:
351
+ av_model.eval()
352
+ cos_vals = []
353
+ with torch.no_grad():
354
+ eval_idx = random.sample(range(len(acts)), min(50, len(acts)))
355
+ for i in eval_idx:
356
+ act_s = scale_acts(acts[i:i+1], dev)
357
+ act_n = acts_n[i:i+1].to(dev)
358
+ expls = av_generate(av_model, tok_av, act_s, dev, args.max_new, temperature=0)
359
+ recon = ar_reconstruct(ar_backbone, ar_head, tok_ar, expls[0], dev)
360
+ cos_vals.append((recon * act_n).sum(-1).item())
361
+ val_cos = sum(cos_vals)/len(cos_vals)
362
+ delta = val_cos - sft_cos
363
+ flag = "✓ BEST" if val_cos > best_cos else ("⚠️ regress" if val_cos < sft_cos - 0.02 else "")
364
+ log(f"[EVAL step {step}] roundtrip_cos={val_cos:.4f} | ΔSFT={delta:+.4f} | best={best_cos:.4f} {flag}")
365
+ writer.add_scalar("rl/val_roundtrip_cosine", val_cos, step)
366
+ writer.add_scalar("rl/delta_vs_sft", delta, step)
367
+
368
+ if val_cos > best_cos:
369
+ best_cos, best_step = val_cos, step
370
+ av_model.save_pretrained(ARTIFACTS / "checkpoints" / "av_rl_best")
371
+ log(f" Saved best RL checkpoint (cos={best_cos:.4f})")
372
+
373
+ # Sample explanation
374
+ sample_act = scale_acts(acts[eval_idx[0]:eval_idx[0]+1], dev)
375
+ expls = av_generate(av_model, tok_av, sample_act, dev, args.max_new, temperature=0)
376
+ writer.add_text("rl/sample_expl", expls[0][0], step)
377
+ log(f" Sample: {expls[0][0][:80]}")
378
+
379
+ if step % args.save_every == 0:
380
+ av_model.save_pretrained(ARTIFACTS / "checkpoints" / f"av_rl_step{step}")
381
+ # Save resume checkpoint
382
+ rl_resume_dir = ARTIFACTS / "checkpoints" / "av_rl_resume"
383
+ rl_resume_dir.mkdir(parents=True, exist_ok=True)
384
+ av_model.save_pretrained(rl_resume_dir)
385
+ torch.save({
386
+ "step": step,
387
+ "optimizer": opt.state_dict(),
388
+ "best_cos": best_cos,
389
+ "best_step": best_step,
390
+ "sft_cos": sft_cos,
391
+ }, rl_resume_dir / "training_state.pt")
392
+ log(f" Saved resume checkpoint at step {step}")
393
+
394
+ log(f"\nRL done | best_cos={best_cos:.4f} at step {best_step} | SFT_cos={sft_cos:.4f}")
395
+ log(f"Best checkpoint: {ARTIFACTS/'checkpoints'/'av_rl_best'}")
396
+ writer.add_hparams(
397
+ {"lr":args.lr_av,"G":args.group_size,"B":args.batch,"kl":args.kl_coef,"temp":args.temperature},
398
+ {"hparam/best_rl_cosine": best_cos, "hparam/sft_cosine": sft_cos}
399
+ )
400
+ writer.close()
401
+
402
+
403
+ if __name__ == "__main__":
404
+ main()
experiments/tiny_nla/train_sft_v2.py ADDED
@@ -0,0 +1,451 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/usr/bin/env python3
2
+ """
3
+ Phase 2: AR + AV SFT on v2 data (9810 records).
4
+ Real-time monitoring via TensorBoard + console.
5
+
6
+ Usage:
7
+ python train_sft_v2.py --stage ar # AR only
8
+ python train_sft_v2.py --stage av # AV only
9
+ python train_sft_v2.py # both (AR first)
10
+
11
+ Monitor:
12
+ tensorboard --logdir artifacts/tiny_nla/runs
13
+ """
14
+ import json, yaml, random, time, argparse, sys, os
15
+ from pathlib import Path
16
+ from datetime import datetime
17
+
18
+ # Disable HuggingFace network access — use local cache only, no HEAD checks to huggingface.co
19
+ os.environ.setdefault("HF_HUB_OFFLINE", "1")
20
+ os.environ.setdefault("TRANSFORMERS_OFFLINE", "1")
21
+ os.environ.setdefault("HF_DATASETS_OFFLINE", "1")
22
+
23
+ import torch
24
+ import torch.nn as nn
25
+ import torch.nn.functional as F
26
+ import pyarrow.parquet as pq
27
+ from torch.utils.data import Dataset, DataLoader
28
+ from torch.utils.tensorboard import SummaryWriter
29
+ from transformers import AutoModelForCausalLM, AutoTokenizer
30
+ from peft import LoraConfig, get_peft_model, PeftModel, TaskType
31
+
32
+ REPO_ROOT = Path(__file__).resolve().parents[2]
33
+ ARTIFACTS = REPO_ROOT / "artifacts" / "tiny_nla"
34
+ META_PATH = Path(__file__).resolve().parent / "nla_meta.yaml"
35
+ RUNS_DIR = ARTIFACTS / "runs"
36
+
37
+ meta = yaml.safe_load(open(META_PATH))
38
+ D_MODEL = meta["d_model"]
39
+ INJ_CHAR = meta["tokens"]["injection_char"]
40
+ INJ_TOK_ID = meta["tokens"]["injection_token_id"]
41
+ INJ_SCALE = meta["extraction"]["injection_scale"]
42
+ BASE_MODEL = meta["base_model"]
43
+ INST_MODEL = meta["av_init_model"]
44
+
45
+
46
+ def get_device():
47
+ if torch.cuda.is_available(): return torch.device("cuda")
48
+ if torch.backends.mps.is_available(): return torch.device("mps")
49
+ return torch.device("cpu")
50
+
51
+
52
+ def log(msg):
53
+ print(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}", flush=True)
54
+
55
+
56
+ # ── Data loading ─────────────────────────────────────
57
+ def load_data():
58
+ table = pq.read_table(ARTIFACTS / "activations_v2.parquet")
59
+ labels = {(r["text_idx"], r["pos"]): r["teacher_explanation"]
60
+ for r in json.load(open(ARTIFACTS / "teacher_labels_v2.json"))}
61
+
62
+ records, acts = [], []
63
+ for i in range(len(table)):
64
+ ti, pos = table["text_idx"][i].as_py(), table["pos"][i].as_py()
65
+ expl = labels.get((ti, pos), "")
66
+ if not expl or len(expl) < 5:
67
+ continue
68
+ act = torch.tensor(table["activation"][i].as_py(), dtype=torch.float32)
69
+ records.append({"text_idx": ti, "pos": pos,
70
+ "text": table["text"][i].as_py(),
71
+ "token_text": table["token_text"][i].as_py(),
72
+ "teacher_explanation": expl})
73
+ acts.append(act)
74
+
75
+ acts_t = torch.stack(acts)
76
+ acts_n = F.normalize(acts_t, dim=-1)
77
+ log(f"Loaded {len(records)} records | norm {acts_t.norm(dim=-1).min():.1f}–{acts_t.norm(dim=-1).max():.1f}")
78
+ return records, acts_t, acts_n
79
+
80
+
81
+ # ══════════════════════════════════════════════════════
82
+ # AR SFT
83
+ # ══════════════════════════════════════════════════════
84
+ class ARHead(nn.Module):
85
+ def __init__(self, d_model):
86
+ super().__init__()
87
+ self.proj = nn.Linear(d_model, d_model, bias=False)
88
+ def forward(self, h):
89
+ return F.normalize(self.proj(h), dim=-1)
90
+
91
+
92
+ class ARDataset(Dataset):
93
+ def __init__(self, records, acts_norm, tok, max_len=96):
94
+ self.items = []
95
+ for rec, act in zip(records, acts_norm):
96
+ ids = tok(rec["teacher_explanation"], truncation=True,
97
+ max_length=max_len, return_tensors="pt")["input_ids"][0]
98
+ self.items.append((ids, act))
99
+ def __len__(self): return len(self.items)
100
+ def __getitem__(self, i): return self.items[i]
101
+
102
+
103
+ def ar_collate(batch, pad_id):
104
+ ids_list, acts = zip(*batch)
105
+ L = max(x.shape[0] for x in ids_list)
106
+ ids = torch.full((len(batch), L), pad_id, dtype=torch.long)
107
+ mask = torch.zeros(len(batch), L, dtype=torch.long)
108
+ for i, x in enumerate(ids_list):
109
+ ids[i, :len(x)] = x; mask[i, :len(x)] = 1
110
+ return ids, mask, torch.stack(acts)
111
+
112
+
113
+ def train_ar(records, acts_t, acts_n):
114
+ dev = get_device()
115
+ run_name = f"ar_{datetime.now().strftime('%m%d_%H%M')}"
116
+ writer = SummaryWriter(RUNS_DIR / run_name)
117
+ log(f"AR SFT | device={dev} | tb: tensorboard --logdir {RUNS_DIR}")
118
+
119
+ tok = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
120
+ tok.pad_token_id = tok.eos_token_id
121
+
122
+ random.seed(42); idx = list(range(len(records))); random.shuffle(idx)
123
+ val_n = max(200, int(len(idx)*0.1))
124
+ tr_idx, va_idx = idx[val_n:], idx[:val_n]
125
+
126
+ pad = tok.eos_token_id
127
+ tr_dl = DataLoader(ARDataset([records[i] for i in tr_idx], acts_n[tr_idx], tok),
128
+ batch_size=32, shuffle=True, collate_fn=lambda b: ar_collate(b, pad))
129
+ va_dl = DataLoader(ARDataset([records[i] for i in va_idx], acts_n[va_idx], tok),
130
+ batch_size=32, shuffle=False, collate_fn=lambda b: ar_collate(b, pad))
131
+
132
+ log(f"Train {len(tr_idx)}, Val {len(va_idx)}")
133
+
134
+ log(f"Loading {BASE_MODEL}...")
135
+ backbone = AutoModelForCausalLM.from_pretrained(
136
+ BASE_MODEL, trust_remote_code=True, dtype=torch.float32,
137
+ low_cpu_mem_usage=True, attn_implementation="eager").to(dev)
138
+ backbone.eval()
139
+ for p in backbone.parameters(): p.requires_grad_(False)
140
+
141
+ head = ARHead(D_MODEL).to(dev)
142
+ opt = torch.optim.AdamW(head.parameters(), lr=3e-4, weight_decay=0.01)
143
+ sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=25, eta_min=1e-5)
144
+
145
+ # --- checkpoint helper: save best immediately to disk ---
146
+ ckpt_dir = ARTIFACTS / "checkpoints" / "ar_v2"
147
+ ckpt_dir.mkdir(parents=True, exist_ok=True)
148
+ def _save_ar_ckpt(state, cos_val, epoch_num):
149
+ p = ckpt_dir / "ar_head_v2.pt"
150
+ torch.save({"head": state, "d_model": D_MODEL,
151
+ "val_cosine": cos_val, "mean_baseline": mean_cos,
152
+ "epoch": epoch_num}, p)
153
+
154
+ # Baselines
155
+ mean_dir = acts_n.mean(0)
156
+ mean_cos = (acts_n @ mean_dir).mean().item()
157
+ log(f"Mean-direction baseline cosine: {mean_cos:.4f}")
158
+ writer.add_scalar("baseline/mean_cosine", mean_cos, 0)
159
+
160
+ best_cos, best_state, global_step = -1.0, None, 0
161
+ patience, pat_cnt = 6, 0
162
+
163
+ for epoch in range(25):
164
+ head.train()
165
+ ep_loss = []
166
+ for ids, mask, gold in tr_dl:
167
+ ids, mask, gold = ids.to(dev), mask.to(dev), gold.to(dev)
168
+ with torch.no_grad():
169
+ h = backbone(ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
170
+ last = h[torch.arange(len(h)), mask.sum(1)-1]
171
+ pred = head(last)
172
+ loss = (2*(1-(pred*gold).sum(-1))).mean()
173
+ opt.zero_grad(); loss.backward(); opt.step()
174
+ ep_loss.append(loss.item())
175
+ writer.add_scalar("ar/train_loss_step", loss.item(), global_step)
176
+ global_step += 1
177
+
178
+ head.eval(); cos_vals = []
179
+ with torch.no_grad():
180
+ for ids, mask, gold in va_dl:
181
+ ids, mask, gold = ids.to(dev), mask.to(dev), gold.to(dev)
182
+ h = backbone(ids, attention_mask=mask, output_hidden_states=True).hidden_states[-1]
183
+ last = h[torch.arange(len(h)), mask.sum(1)-1]
184
+ pred = head(last)
185
+ cos_vals.extend((pred*gold).sum(-1).tolist())
186
+
187
+ train_loss = sum(ep_loss)/len(ep_loss)
188
+ val_cos = sum(cos_vals)/len(cos_vals)
189
+ sched.step()
190
+
191
+ writer.add_scalar("ar/train_loss", train_loss, epoch)
192
+ writer.add_scalar("ar/val_cosine", val_cos, epoch)
193
+ writer.add_scalar("ar/lr", opt.param_groups[0]["lr"], epoch)
194
+
195
+ delta = val_cos - mean_cos
196
+ flag = "✓" if val_cos > best_cos else "↓"
197
+ log(f"AR Epoch {epoch+1:2d} | train_loss={train_loss:.4f} | val_cos={val_cos:.4f} "
198
+ f"| Δmean={delta:+.4f} | best={best_cos:.4f} {flag}")
199
+
200
+ if val_cos > best_cos:
201
+ best_cos, pat_cnt = val_cos, 0
202
+ best_state = {k: v.clone() for k, v in head.state_dict().items()}
203
+ _save_ar_ckpt(best_state, best_cos, epoch + 1) # flush to disk NOW
204
+ else:
205
+ pat_cnt += 1
206
+ if pat_cnt >= patience:
207
+ log(f"Early stop (patience={patience})")
208
+ break
209
+
210
+ # --- safety net: ensure best state is on disk even on interrupt ---
211
+ head.load_state_dict(best_state)
212
+ _save_ar_ckpt(best_state, best_cos, epoch + 1)
213
+
214
+ writer.add_hparams({"lr": 3e-4, "batch": 32, "d_model": D_MODEL},
215
+ {"hparam/best_val_cosine": best_cos})
216
+ writer.close()
217
+ log(f"AR done | best_cos={best_cos:.4f} vs mean={mean_cos:.4f} Δ={best_cos-mean_cos:+.4f}")
218
+ log(f"Checkpoint: {ckpt_dir / 'ar_head_v2.pt'}")
219
+ return best_cos
220
+
221
+
222
+ # ══════════════════════════════════════════════════════
223
+ # AV SFT
224
+ # ══════════════════════════════════════════════════════
225
+ class AVDataset(Dataset):
226
+ def __init__(self, records, acts, tok, max_expl=80):
227
+ prompt_ids = tok(f"<concept>{INJ_CHAR}</concept>\n<explanation>",
228
+ return_tensors="pt")["input_ids"][0]
229
+ self.plen = len(prompt_ids)
230
+ self.items = []
231
+ for rec, act in zip(records, acts):
232
+ expl_ids = tok(rec["teacher_explanation"], add_special_tokens=False,
233
+ return_tensors="pt")["input_ids"][0][:max_expl]
234
+ eos = torch.tensor([tok.eos_token_id])
235
+ ids = torch.cat([prompt_ids, expl_ids, eos])
236
+ lbl = ids.clone(); lbl[:self.plen] = -100
237
+ self.items.append({"input_ids": ids, "labels": lbl, "act": act,
238
+ "teacher": rec["teacher_explanation"]})
239
+ def __len__(self): return len(self.items)
240
+ def __getitem__(self, i): return self.items[i]
241
+
242
+
243
+ def av_collate(batch, pad_id):
244
+ L = max(b["input_ids"].shape[0] for b in batch)
245
+ ids = torch.full((len(batch), L), pad_id, dtype=torch.long)
246
+ mask = torch.zeros(len(batch), L, dtype=torch.long)
247
+ lbl = torch.full((len(batch), L), -100, dtype=torch.long)
248
+ acts = torch.stack([b["act"] for b in batch])
249
+ for i, b in enumerate(batch):
250
+ n = b["input_ids"].shape[0]
251
+ ids[i,:n] = b["input_ids"]; mask[i,:n] = 1; lbl[i,:n] = b["labels"]
252
+ return {"ids": ids, "mask": mask, "lbl": lbl, "acts": acts}
253
+
254
+
255
+ def scale_acts(acts, dev):
256
+ a = acts.to(dev).float()
257
+ return a / a.norm(dim=-1, keepdim=True).clamp(1e-6) * INJ_SCALE
258
+
259
+
260
+ class AVModel(nn.Module):
261
+ def __init__(self, lora_model):
262
+ super().__init__(); self.model = lora_model
263
+ def forward(self, ids, mask, lbl, acts):
264
+ e = self.model.get_input_embeddings()(ids)
265
+ # Vectorized activation injection (no Python loop)
266
+ positions = (ids == INJ_TOK_ID).float().argmax(dim=1) # [B]
267
+ b_idx = torch.arange(ids.shape[0], device=ids.device)
268
+ e[b_idx, positions] = acts.to(e.dtype)
269
+ return self.model(inputs_embeds=e, attention_mask=mask, labels=lbl)
270
+
271
+
272
+ def train_av(records, acts_t, resume=False, batch_size=8, lr=5e-5, epochs=30):
273
+ dev = get_device()
274
+ run_name = f"av_{datetime.now().strftime('%m%d_%H%M')}"
275
+ writer = SummaryWriter(RUNS_DIR / run_name)
276
+ log(f"AV SFT | device={dev} | tb: tensorboard --logdir {RUNS_DIR}")
277
+
278
+ tok = AutoTokenizer.from_pretrained(INST_MODEL, trust_remote_code=True)
279
+ tok.pad_token_id = tok.eos_token_id
280
+
281
+ random.seed(42); idx = list(range(len(records))); random.shuffle(idx)
282
+ val_n = max(200, int(len(idx)*0.1))
283
+ tr_idx, va_idx = idx[val_n:], idx[:val_n]
284
+
285
+ pad = tok.eos_token_id
286
+ tr_ds = AVDataset([records[i] for i in tr_idx], acts_t[tr_idx], tok)
287
+ va_ds = AVDataset([records[i] for i in va_idx], acts_t[va_idx], tok)
288
+ tr_dl = DataLoader(tr_ds, batch_size=batch_size, shuffle=True,
289
+ collate_fn=lambda b: av_collate(b, pad))
290
+ va_dl = DataLoader(va_ds, batch_size=batch_size, shuffle=False,
291
+ collate_fn=lambda b: av_collate(b, pad))
292
+ log(f"Train {len(tr_idx)}, Val {len(va_idx)}")
293
+
294
+ log(f"Loading {INST_MODEL} (float16, sdpa)...")
295
+ base = AutoModelForCausalLM.from_pretrained(
296
+ INST_MODEL, trust_remote_code=True, dtype=torch.float16,
297
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(dev)
298
+ resume_dir = ARTIFACTS / "checkpoints" / "av_v2_resume"
299
+ av_v2_dir = ARTIFACTS / "checkpoints" / "av_v2"
300
+ if resume and (resume_dir / "training_state.pt").exists():
301
+ lora_model = PeftModel.from_pretrained(base, resume_dir, is_trainable=True)
302
+ log(f"Resumed LoRA + optimizer from {resume_dir}")
303
+ elif resume and av_v2_dir.exists():
304
+ lora_model = PeftModel.from_pretrained(base, av_v2_dir, is_trainable=True)
305
+ log(f"⚠️ No resume checkpoint, loading best SFT from {av_v2_dir} (optimizer reset, epoch=0)")
306
+ else:
307
+ if resume:
308
+ log(f"⚠️ No checkpoint found, starting fresh LoRA")
309
+ lora_model = get_peft_model(base, LoraConfig(
310
+ task_type=TaskType.CAUSAL_LM, r=8, lora_alpha=16,
311
+ lora_dropout=0.1, target_modules=["q_proj","v_proj"], bias="none"))
312
+ lora_model.print_trainable_parameters()
313
+
314
+ av = AVModel(lora_model).to(dev)
315
+ opt = torch.optim.AdamW(av.parameters(), lr=lr, weight_decay=0.05)
316
+ sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=epochs, eta_min=lr/10)
317
+
318
+ best_loss, pat_cnt, global_step = float("inf"), 0, 0
319
+ start_epoch = 0
320
+ PATIENCE = 8
321
+
322
+ if resume and (resume_dir / "training_state.pt").exists():
323
+ ts = torch.load(resume_dir / "training_state.pt",
324
+ map_location=dev, weights_only=False)
325
+ opt.load_state_dict(ts["optimizer"])
326
+ sched.load_state_dict(ts["scheduler"])
327
+ start_epoch = ts["epoch"] + 1
328
+ best_loss = ts["best_loss"]
329
+ pat_cnt = ts["pat_cnt"]
330
+ global_step = ts["global_step"]
331
+ log(f"Resumed from epoch {start_epoch} | best_loss={best_loss:.4f} | pat_cnt={pat_cnt}")
332
+
333
+ for epoch in range(start_epoch, epochs):
334
+ av.train(); ep_loss = []
335
+ for batch in tr_dl:
336
+ ids = batch["ids"].to(dev); mask = batch["mask"].to(dev)
337
+ lbl = batch["lbl"].to(dev); acts = scale_acts(batch["acts"], dev)
338
+ out = av(ids, mask, lbl, acts)
339
+ opt.zero_grad(); out.loss.backward()
340
+ torch.nn.utils.clip_grad_norm_(av.parameters(), 1.0)
341
+ opt.step()
342
+ ep_loss.append(out.loss.item())
343
+ writer.add_scalar("av/train_loss_step", out.loss.item(), global_step)
344
+ global_step += 1
345
+
346
+ av.eval(); vl = []
347
+ with torch.no_grad():
348
+ for batch in va_dl:
349
+ out = av(batch["ids"].to(dev), batch["mask"].to(dev),
350
+ batch["lbl"].to(dev), scale_acts(batch["acts"], dev))
351
+ vl.append(out.loss.item())
352
+
353
+ train_loss = sum(ep_loss)/len(ep_loss)
354
+ val_loss = sum(vl)/len(vl)
355
+ sched.step()
356
+
357
+ writer.add_scalar("av/train_loss", train_loss, epoch)
358
+ writer.add_scalar("av/val_loss", val_loss, epoch)
359
+ writer.add_scalar("av/lr", opt.param_groups[0]["lr"], epoch)
360
+
361
+ # overfit alert
362
+ gap = train_loss / max(val_loss, 1e-6)
363
+ flag = "⚠️ overfit" if gap < 0.3 else ("✓" if val_loss < best_loss else "↓")
364
+ log(f"AV Epoch {epoch+1:2d} | train={train_loss:.4f} | val={val_loss:.4f} "
365
+ f"| gap={gap:.2f} | best={best_loss:.4f} {flag}")
366
+
367
+ if val_loss < best_loss:
368
+ best_loss, pat_cnt = val_loss, 0
369
+ av.model.save_pretrained(ARTIFACTS / "checkpoints" / "av_v2")
370
+ tok.save_pretrained(ARTIFACTS / "checkpoints" / "av_v2")
371
+ else:
372
+ pat_cnt += 1
373
+
374
+ # Save resume checkpoint every epoch
375
+ resume_dir = ARTIFACTS / "checkpoints" / "av_v2_resume"
376
+ resume_dir.mkdir(parents=True, exist_ok=True)
377
+ av.model.save_pretrained(resume_dir)
378
+ tok.save_pretrained(resume_dir)
379
+ torch.save({
380
+ "epoch": epoch,
381
+ "optimizer": opt.state_dict(),
382
+ "scheduler": sched.state_dict(),
383
+ "best_loss": best_loss,
384
+ "pat_cnt": pat_cnt,
385
+ "global_step": global_step,
386
+ }, resume_dir / "training_state.pt")
387
+
388
+ if pat_cnt >= PATIENCE:
389
+ log(f"Early stop (patience={PATIENCE})"); break
390
+
391
+ # Sample generations from best checkpoint
392
+ log("Generating samples from best checkpoint...")
393
+ base2 = AutoModelForCausalLM.from_pretrained(
394
+ INST_MODEL, trust_remote_code=True, dtype=torch.float16,
395
+ low_cpu_mem_usage=True, attn_implementation="sdpa").to(dev)
396
+ best_av = PeftModel.from_pretrained(base2, ARTIFACTS/"checkpoints"/"av_v2").to(dev)
397
+ best_av.eval()
398
+
399
+ prompt = f"<concept>{INJ_CHAR}</concept>\n<explanation>"
400
+ p_ids = tok(prompt, return_tensors="pt")["input_ids"].to(dev)
401
+ p_mask = torch.ones(1, p_ids.shape[1], device=dev)
402
+ nonempty = 0
403
+ with torch.no_grad():
404
+ for i in range(min(10, len(va_ds))):
405
+ act = va_ds[i]["act"].unsqueeze(0)
406
+ sc = scale_acts(act, dev)
407
+ emb = best_av.get_input_embeddings()(p_ids)
408
+ pos = (p_ids[0] == INJ_TOK_ID).nonzero(as_tuple=True)[0]
409
+ if len(pos): emb[0, pos[0]] = sc[0].to(emb.dtype)
410
+ out = best_av.generate(inputs_embeds=emb, attention_mask=p_mask,
411
+ max_new_tokens=80, do_sample=False,
412
+ pad_token_id=tok.eos_token_id)
413
+ gen = tok.decode(out[0], skip_special_tokens=True).strip()
414
+ if gen: nonempty += 1
415
+ writer.add_text("av/samples", f"T: {va_ds[i]['teacher'][:60]}\nG: {gen[:80]}", i)
416
+ log(f" [{i}] T: {va_ds[i]['teacher'][:55]}")
417
+ log(f" G: {gen[:80]}")
418
+
419
+ log(f"Non-empty: {nonempty}/10")
420
+ writer.add_hparams({"lr":lr,"lora_r":8,"batch":batch_size},
421
+ {"hparam/best_val_loss": best_loss})
422
+ writer.close()
423
+ log(f"AV done | best_val_loss={best_loss:.4f}")
424
+ log(f"Checkpoint: {ARTIFACTS/'checkpoints'/'av_v2'}")
425
+ return best_loss
426
+
427
+
428
+ def main():
429
+ parser = argparse.ArgumentParser()
430
+ parser.add_argument("--stage", choices=["ar","av","both"], default="both")
431
+ parser.add_argument("--resume", action="store_true",
432
+ help="Resume AV training from latest checkpoint")
433
+ parser.add_argument("--batch-size", type=int, default=8,
434
+ help="Batch size (smaller = more SGD noise = better generalization)")
435
+ parser.add_argument("--lr", type=float, default=5e-5,
436
+ help="Peak LR (lower for warm restart from existing checkpoint)")
437
+ parser.add_argument("--epochs", type=int, default=30)
438
+ args = parser.parse_args()
439
+
440
+ RUNS_DIR.mkdir(parents=True, exist_ok=True)
441
+ records, acts_t, acts_n = load_data()
442
+
443
+ if args.stage in ("ar","both"):
444
+ train_ar(records, acts_t, acts_n)
445
+ if args.stage in ("av","both"):
446
+ train_av(records, acts_t, resume=args.resume,
447
+ batch_size=args.batch_size, lr=args.lr, epochs=args.epochs)
448
+
449
+
450
+ if __name__ == "__main__":
451
+ main()
requirements.txt CHANGED
@@ -1,6 +1,7 @@
1
  transformers>=4.51.0,<5.0.0 # 5.x 需 accelerate库,HF Spaces 上易出问题
2
  torch>=2.1.0,<2.5.0
3
  numpy>=1.24.0
 
4
  connexion[flask,swagger-ui,uvicorn]>=3.0.0
5
  flask>=3.0.0
6
  PyYAML>=6.0
 
1
  transformers>=4.51.0,<5.0.0 # 5.x 需 accelerate库,HF Spaces 上易出问题
2
  torch>=2.1.0,<2.5.0
3
  numpy>=1.24.0
4
+ peft>=0.14.0 # LoRA 模型加载(Tiny-NLA Phase 1b)
5
  connexion[flask,swagger-ui,uvicorn]>=3.0.0
6
  flask>=3.0.0
7
  PyYAML>=6.0
server.py CHANGED
@@ -40,6 +40,7 @@ from backend.api.analyze_semantic import analyze_semantic # noqa: F401
40
  from backend.api.prediction_attribute import prediction_attribute # noqa: F401
41
  from backend.api.ablation_attribute import ablation_attribute # noqa: F401
42
  from backend.api.logit_lens import logit_lens # noqa: F401
 
43
  from backend.api.branch_next import branch_next # noqa: F401
44
  from backend.api.tokenize import tokenize # noqa: F401
45
  from backend.api.model_switch import ( # noqa: F401
 
40
  from backend.api.prediction_attribute import prediction_attribute # noqa: F401
41
  from backend.api.ablation_attribute import ablation_attribute # noqa: F401
42
  from backend.api.logit_lens import logit_lens # noqa: F401
43
+ from backend.api.activation_explain import activation_explain # noqa: F401
44
  from backend.api.branch_next import branch_next # noqa: F401
45
  from backend.api.tokenize import tokenize # noqa: F401
46
  from backend.api.model_switch import ( # noqa: F401
server.yaml CHANGED
@@ -763,6 +763,67 @@ paths:
763
  503:
764
  description: 服务繁忙
765
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
766
  /branch-next:
767
  post:
768
  tags:
 
763
  503:
764
  description: 服务繁忙
765
 
766
+ /activation-explain:
767
+ post:
768
+ tags:
769
+ - all
770
+ summary: activation explain (Tiny-NLA)
771
+ operationId: server.activation_explain
772
+ parameters:
773
+ - in: body
774
+ name: activation_explain_request
775
+ schema:
776
+ type: object
777
+ required:
778
+ - model
779
+ - source_page
780
+ properties:
781
+ model:
782
+ type: string
783
+ enum: [base, instruct]
784
+ source_page:
785
+ type: string
786
+ text:
787
+ type: string
788
+ description: 输入上下文文本(与 token_index 配套)
789
+ token_index:
790
+ type: integer
791
+ description: 目标 token 索引(text 模式必填)
792
+ vector:
793
+ type: array
794
+ items:
795
+ type: number
796
+ description: 1024 维激活向量(与 text 互斥)
797
+ responses:
798
+ 200:
799
+ description: 激活解释结果
800
+ schema:
801
+ type: object
802
+ properties:
803
+ success:
804
+ type: boolean
805
+ concept:
806
+ type: string
807
+ description: 激活概念标签
808
+ explanation:
809
+ type: string
810
+ description: 自然语言解释
811
+ roundtrip_cosine:
812
+ type: number
813
+ description: roundtrip 可信度(0-1)
814
+ vector_dim:
815
+ type: integer
816
+ description: 向量维度
817
+ note:
818
+ type: string
819
+ description: 备注或状态说明
820
+ 400:
821
+ description: 缺少必要字段、model 非法、text/vector 均未提供
822
+ 500:
823
+ description: 推理失败
824
+ 503:
825
+ description: 服务繁忙
826
+
827
  /branch-next:
828
  post:
829
  tags:
train ADDED
@@ -0,0 +1,394 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ 下面是更完整版,可直接丢给训练 agent。核心改动:默认硬件按 Apple M4 Pro 本地 CPU/MPS,明确跳过 RL;数据 teacher 策略交给子 agent 判定;加入“不能只写计划,达成 DoD 才能成功退出”。
3
+
4
+ ```text
5
+ # 任务:在 InfoLens 上训练一个 Qwen3-0.6B Tiny-NLA(AV + AR)
6
+
7
+ 你是训练执行 agent。你的任务不是写调研报告,而是把一个能跑的 Tiny-NLA 原型训练出来,并交付可复现脚本、checkpoint、评估样例。
8
+
9
+ 请完整读完本任务书再动手。任何红线冲突立即停止并报告。除非遇到明确硬阻塞,否则你不能只输出计划后退出;你必须持续执行、修正、评估,直到满足“Definition of Done”。
10
+
11
+ ---
12
+
13
+ ## 0. 已知项目与硬件约束
14
+
15
+ Repo: `/Users/cccmmd/InfoLens`
16
+
17
+ InfoLens 是本地 LLM 可解释性工具,已有:
18
+ - `/api/prediction-attribute`:next-token attribution
19
+ - `/api/ablation-attribute`:ablation attribution
20
+ - `/api/logit-lens`:逐层 hidden state 经过 final norm + lm_head 的 top-k 与目标 token 概率轨迹
21
+
22
+ 默认模型见 `model_paths.py`:
23
+ - Base: `qwen3-0.6b` -> `Qwen/Qwen3-0.6B-Base`
24
+ - Instruct: `qwen3-0.6b-instruct` -> `Qwen/Qwen3-0.6B`
25
+
26
+ 硬件已知:
27
+ - 本机是 Apple M4 Pro 本地机器。
28
+ - 不要假设有 CUDA。
29
+ - 可以探测 MPS,但必须先 smoke test 反向传播;MPS 不可靠时退回 CPU。
30
+ - 默认策略:只做 SFT,不做 RL。
31
+ - RL / GRPO 在本机视为 out of scope,除非用户另行提供远程 CUDA 训练机。
32
+
33
+ ---
34
+
35
+ ## 1. 红线
36
+
37
+ 违反任意一条即任务失败,立即停止并报告:
38
+
39
+ 1. 禁止下载或运行 7B 及以上模型。
40
+ 2. 禁止使用官方 released NLA checkpoint。它们绑定 Qwen2.5-7B / Gemma / Llama 的激活空间,与 Qwen3-0.6B 不兼容。
41
+ 3. 禁止照搬 Miles + SGLang + Megatron 训练栈。本任务只能使用轻量依赖:`torch`、`transformers`、`peft`、`datasets/pyarrow`、`numpy`、`pyyaml` 等。
42
+ 4. 禁止把官方 7B/70B 超参当成默认值。0.6B 必须自己做小规模 smoke 与轻量调参。
43
+ 5. 禁止把大模型权重、parquet 数据、训练 artifacts 提交到 git。
44
+ 6. 不要改动生产 UI/API,除非用户后续明确要求。本任务只做实验脚本与 artifacts。
45
+
46
+ ---
47
+
48
+ ## 2. 目标
49
+
50
+ 训练一对 Tiny-NLA 组件,用于解释 `Qwen/Qwen3-0.6B-Base` 某一层 residual stream activation。
51
+
52
+ - AV / Activation Verbalizer: `activation vector -> natural language explanation`
53
+ - AR / Activation Reconstructor: `explanation -> reconstructed activation vector`
54
+ - 比较向量前必须 L2 normalize。
55
+ - round-trip loss 使用 `MSE = 2 * (1 - cosine)` 或等价 normalized MSE。
56
+ - 最终 InfoLens 更依赖 AV:输入某层 activation,输出一句可读中文解释。
57
+ - AR 用于客观评估与未来 reward,不要求达到官方 7B 水平,但必须训练、评估、和 baseline 比较。
58
+
59
+ ---
60
+
61
+ ## 3. 成功退出条件 Definition of Done
62
+
63
+ 你不能在满足以下条件前声称任务完成:
64
+
65
+ 1. 已确认并记录环境:
66
+ - device: CPU / MPS / CUDA
67
+ - `Qwen/Qwen3-0.6B-Base` 的 `num_hidden_layers` 与 `hidden_size`
68
+ - 选择的 layer index,按约 2/3 深度计算,并说明理由
69
+ - injection token 是否为单 token
70
+ - injection scale 如何估计
71
+
72
+ 2. 已完成可复现数据生成:
73
+ - 至少 smoke 数据 200 条
74
+ - 如果速度允许,扩到 500-2000 条
75
+ - 数据包含:context、token index、layer、activation_vector、teacher explanation、target token/top-k debug 信息
76
+ - 数据与 sidecar 存在 artifacts 目录,且不进入 git
77
+
78
+ 3. 已完成 Stage 0 smoke:
79
+ - 能提取 Qwen3-0.6B-Base hidden state
80
+ - 能用 `input_embeds` 注入 activation
81
+ - 能跑一次 AV forward/generation,不崩溃、不 shape mismatch
82
+
83
+ 4. 已完成 AR SFT:
84
+ - 有训练脚本
85
+ - 有 checkpoint
86
+ - 有 val metrics
87
+ - 必须和 mean baseline / shuffled baseline 对比
88
+ - 如果 AR 训练失败,必须至少做两轮合理修正后才能报告 blocker
89
+
90
+ 5. 已完成 AV SFT:
91
+ - 有训练脚本
92
+ - 有 checkpoint 或 LoRA adapter
93
+ - 能对 held-out activation 生成中文解释
94
+ - 至少输出 20 条 worked examples
95
+ - 20 条里不能大面积乱码、空输出、模板废话;若质量很差,必须继续修正数据或训练设置,不能直接交付
96
+
97
+ 6. 已交付推理脚本:
98
+ - 输入文本 + token position,自动提取 selected layer activation
99
+ - 调用 AV 输出解释
100
+ - 若 AR 可用,同时输出 reconstruction cosine/MSE
101
+ - 从 `nla_meta.yaml` 读取配置,不硬编码 layer/token/scale/template
102
+
103
+ 7. 已交付最终报告:
104
+ - 环境
105
+ - 数据策略
106
+ - 训练耗时
107
+ - AR 指标
108
+ - AV 质量观察
109
+ - 20 条样例
110
+ - 失败案例与局限
111
+ - 下一步建议
112
+
113
+ 只有以上完成,才能输出“任务完成”。否则只能输出“阻塞报告”,并附证据与已尝试修复项。
114
+
115
+ ---
116
+
117
+ ## 4. 外部参考,只学算法,不照搬基础设施
118
+
119
+ 参考仓库:
120
+ `https://github.com/kitft/natural_language_autoencoders`
121
+
122
+ 开工前阅读:
123
+ - `README.md`
124
+ - `docs/inference.md`
125
+ - `docs/design.md`
126
+ - `nla/schema.py`
127
+ - `nla/config.py`
128
+ - `nla/models.py`
129
+ - `nla/loss.py`
130
+ - `nla/reward.py`
131
+ - `nla/datagen/`
132
+
133
+ 你要借鉴:
134
+ - AV 把 activation 当作一个虚拟 token embedding 注入 prompt
135
+ - AR 从 explanation text 重建 activation
136
+ - sidecar 记录 prompt、injection token、layer、scale、d_model
137
+ - normalized MSE / cosine 作为评估
138
+
139
+ 你不要借鉴:
140
+ - 7B+ released checkpoints
141
+ - Miles / SGLang / Megatron
142
+ - 多 H100 训练配置
143
+ - RL 默认流程
144
+
145
+ ---
146
+
147
+ ## 5. 数据生成策略:必须交给子 agent 判断与执行建议
148
+
149
+ 在生成 teacher explanation 前,先启动一个 focused subagent,任务是:
150
+
151
+ “判断当前环境是否可使用 Claude/外部 API 生成 Tiny-NLA teacher explanations;如果可用,给出低预算批量生成策略;如果不可用,给出本地 `Qwen/Qwen3-0.6B` instruct 生成策略。必须返回具体 prompt 模板、批大小、成本/速度风险、fallback 方案。”
152
+
153
+ 子 agent 必须检查:
154
+ - 是否存在 `ANTHROPIC_API_KEY`
155
+ - 是否存在其他可用 teacher API key
156
+ - 用户是否已明确预算
157
+ - 如果无法确认预算,不要擅自大规模调用外部 API
158
+
159
+ 主 agent 根据子 agent 结论执行:
160
+
161
+ ### 有 Claude API 且预算明确
162
+ - 先生成 200 条 smoke teacher explanations
163
+ - 人工/程序抽查质量
164
+ - 再扩到 500-2000 条
165
+ - 每条 explanation 优先中文,短句,描述该位置模型可能关注的语义
166
+
167
+ ### 没有 API 或预算不明确
168
+ - 使用本地 `Qwen/Qwen3-0.6B` instruct 做 weak teacher
169
+ - 允许降级,用户接受本地 teacher 质量较差
170
+ - 必须在报告中标注:teacher 是 weak local teacher,不是真正 Claude-quality NLA labels
171
+
172
+ Teacher prompt 应包含:
173
+ - 原始 context
174
+ - token 位置
175
+ - target token / final top-k
176
+ - logit lens 中该层附近的 top-k 摘要(如果容易取得)
177
+ - 要求输出 1-2 句中文解释,不要长篇推理
178
+
179
+ 注意:AR 的标签始终是原始 activation vector,不依赖 teacher API。
180
+
181
+ ---
182
+
183
+ ## 6. 推荐目录结构
184
+
185
+ 把实验放在独立目录,例如:
186
+
187
+ `experiments/tiny_nla/`
188
+
189
+ 建议文件:
190
+ - `extract_activations.py`
191
+ - `generate_teacher_labels.py`
192
+ - `train_ar.py`
193
+ - `train_av.py`
194
+ - `eval_roundtrip.py`
195
+ - `infer_tiny_nla.py`
196
+ - `sidecar.py`
197
+ - `README.md`
198
+
199
+ Artifacts 放到:
200
+ - `artifacts/tiny_nla/...`
201
+
202
+ 如果 artifacts 目录未被 gitignore,先加入 gitignore。不要提交模型权重或数据。
203
+
204
+ ---
205
+
206
+ ## 7. Stage 0:环境与注入 smoke
207
+
208
+ 先写并运行 smoke,不要直接训练。
209
+
210
+ 必须做:
211
+ 1. 加载 `Qwen/Qwen3-0.6B-Base`
212
+ 2. 读取真实:
213
+ - `num_hidden_layers`
214
+ - `hidden_size`
215
+ - vocab size
216
+ 3. 计算:
217
+ - `layer_index = round(num_hidden_layers * 2 / 3)`
218
+ 4. 对几条文本跑:
219
+ - `output_hidden_states=True`
220
+ - 取 selected layer 的最后 token 或多个 token activation
221
+ 5. 统计 activation L2 norm:
222
+ - mean
223
+ - p50
224
+ - p90
225
+ - max
226
+ 6. 选择 `injection_scale`:
227
+ - 初始用 p50 或 mean
228
+ - 写入 sidecar
229
+ 7. 选择 injection char:
230
+ - 必须是 tokenizer 下单 token
231
+ - 例如先测试 `㈎`,不行就找其他 rare single token
232
+ 8. 构造 prompt template:
233
+ - `<concept>{injection_char}</concept>`
234
+ - 要求输出 `<explanation>...</explanation>`
235
+ 9. 使用 `input_embeds` 替换 injection token embedding,跑一次 forward/generation。
236
+
237
+ Stage 0 没过,不准训练。
238
+
239
+ ---
240
+
241
+ ## 8. Stage 1:AR SFT
242
+
243
+ AR 输入 explanation text,输出 activation vector。
244
+
245
+ 实现要求:
246
+ - 初版可以用 `Qwen/Qwen3-0.6B-Base` 或 instruct trunk
247
+ - 优先冻结大部分模型,只训练轻量 head 或 LoRA + head
248
+ - AR head: `Linear(d_model, d_model)`
249
+ - 取最后一个 token hidden state 过 head
250
+ - pred 和 gold 都 normalize 后算 MSE
251
+ - 训练集/验证集拆分固定 seed
252
+
253
+ 必须评估:
254
+ - val cosine mean
255
+ - val normalized MSE
256
+ - mean-vector baseline
257
+ - shuffled-label baseline
258
+ - 至少保存 best checkpoint
259
+
260
+ 如果 AR 不明显超过 baseline:
261
+ - 尝试至少两项修正:
262
+ - 改 learning rate
263
+ - 改 batch size
264
+ - 改是否训练 LoRA
265
+ - 清洗 teacher explanation
266
+ - 增加数据量
267
+ - 仍失败再报告,但不要编造成功。
268
+
269
+ ---
270
+
271
+ ## 9. Stage 2:AV SFT
272
+
273
+ AV 输入 activation vector 注入 prompt,输出 teacher explanation。
274
+
275
+ 实现要求:
276
+ - 初始权重优先 `Qwen/Qwen3-0.6B` instruct
277
+ - 使用 LoRA,避免全量微调
278
+ - 通过 `input_embeds` 注入 selected layer activation
279
+ - loss 只算 explanation response token,不算 prompt token
280
+ - prompt template 和 injection 参数全部从 sidecar 读取
281
+ - 输出中文为目标,不是 bug
282
+
283
+ 训练约束:
284
+ - Apple M4 Pro 本地机,不要追求大 batch
285
+ - CPU 慢就降低数据量与 epoch
286
+ - MPS 可用才用 MPS;MPS 出现反向/dtype 问题立即退回 CPU
287
+ - 不要使用 fp16 反向作为默认;优先 fp32/bf16 smoke 后再决定
288
+
289
+ 必须评估:
290
+ - held-out 20 条 worked examples
291
+ - 每��包含:
292
+ - context
293
+ - token text / token index
294
+ - selected layer
295
+ - final top-k
296
+ - teacher explanation
297
+ - AV generated explanation
298
+ - AR reconstruction cosine/MSE(如果 AR 可用)
299
+ - 简短人工判断:相关 / 部分相关 / 不相关
300
+
301
+ 如果 AV 输出乱码、空、完全模板化:
302
+ - 不准交付
303
+ - 必须调整 teacher prompt、训练模板、learning rate、epoch、或数据清洗后重训
304
+
305
+ ---
306
+
307
+ ## 10. Stage 3:RL 明确跳过
308
+
309
+ 本机是 Apple M4 Pro 本地 CPU/MPS 环境,默认跳过 RL。
310
+
311
+ 不要实现 GRPO。
312
+ 不要安装 Miles/SGLang。
313
+ 不要声称完成 RL。
314
+
315
+ 最终报告中写:
316
+ “由于本任务硬件为 Apple M4 Pro 本地 CPU/MPS,无 CUDA 多 GPU,RL/GRPO 阶段按任务约束跳过。本次交付 SFT Tiny-NLA。”
317
+
318
+ ---
319
+
320
+ ## 11. sidecar 契约
321
+
322
+ 必须生成 `nla_meta.yaml`,至少包含:
323
+
324
+ ```yaml
325
+ kind: tiny_nla_model
326
+ base_model: Qwen/Qwen3-0.6B-Base
327
+ av_init_model: Qwen/Qwen3-0.6B
328
+ layer_index: <int>
329
+ num_hidden_layers: <int>
330
+ d_model: <int>
331
+ activation_source: residual_stream
332
+ token_position_policy: selected_token_or_last_token
333
+ extraction:
334
+ injection_scale: <float>
335
+ mse_normalization: l2_direction
336
+ tokens:
337
+ injection_char: "<char>"
338
+ injection_token_id: <int>
339
+ prompt_templates:
340
+ av: |
341
+ ...
342
+ ar: |
343
+ ...
344
+ training:
345
+ device: cpu_or_mps
346
+ dtype: fp32_or_bf16
347
+ dataset_size: <int>
348
+ teacher: claude_or_local_qwen_instruct
349
+ created_at: <timestamp>
350
+ ```
351
+
352
+ 推理脚本必须读 sidecar,不要把这些值散落在代码里。
353
+
354
+ ---
355
+
356
+ ## 12. 最终报告格式
357
+
358
+ 最终报告必须包含:
359
+
360
+ 1. 是否完成 DoD
361
+ 2. 环境报告
362
+ 3. 数据生成策略与 teacher 来源
363
+ 4. 模型与层选择
364
+ 5. injection token 与 scale 统计
365
+ 6. AR 指标与 baseline 对比
366
+ 7. AV 训练设置与质量总结
367
+ 8. 20 条 worked examples 文件路径
368
+ 9. checkpoint / adapter 路径
369
+ 10. 推理脚本用法
370
+ 11. 已知局限
371
+ 12. 下一步建议
372
+
373
+ 不要只说“训练完成”。必须给路径、命令、指标、样例。
374
+
375
+ ---
376
+
377
+ ## 13. 阻塞时如何退出
378
+
379
+ 只有以下情况允许未完成 DoD 而退出:
380
+
381
+ - Qwen3-0.6B 权重无法下载或加载,且重试后失败
382
+ - 本机内存不足,连 Stage 0 smoke 都无法完成
383
+ - tokenizer 找不到合适 single-token injection char,尝试多个候选后失败
384
+ - PyTorch/transformers 在本机无法完成最小 forward/backward,且已给出错误日志
385
+ - 数据 teacher 完全不可用,且本地 instruct 也无法加载
386
+
387
+ 阻塞报告必须包含:
388
+ - 卡在哪个 stage
389
+ - 已尝试哪些修复
390
+ - 完整错误摘要
391
+ - 下一步需要用户提供什么
392
+
393
+ 否则继续工作,直到满足 Definition of Done。
394
+ ```