Spaces:
Running
Running
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)
- .gitignore +4 -0
- DESIGN_PROMPT.md +608 -0
- artifacts/tiny_nla/checkpoints/ar/best_ar_head.pt +3 -0
- artifacts/tiny_nla/checkpoints/av/README.md +207 -0
- artifacts/tiny_nla/checkpoints/av/adapter_config.json +43 -0
- artifacts/tiny_nla/checkpoints/av/adapter_model.safetensors +3 -0
- artifacts/tiny_nla/checkpoints/av/added_tokens.json +28 -0
- artifacts/tiny_nla/checkpoints/av/chat_template.jinja +89 -0
- artifacts/tiny_nla/checkpoints/av/merges.txt +0 -0
- artifacts/tiny_nla/checkpoints/av/special_tokens_map.json +25 -0
- artifacts/tiny_nla/checkpoints/av/tokenizer.json +3 -0
- artifacts/tiny_nla/checkpoints/av/tokenizer_config.json +239 -0
- artifacts/tiny_nla/checkpoints/av/vocab.json +0 -0
- backend/api/activation_explain.py +100 -0
- backend/core/tiny_nla.py +188 -0
- backend/platform/source_page.py +1 -1
- client/src/css/pages/logit_lens.scss +48 -0
- client/src/logit_lens.html +10 -0
- client/src/pages/logit_lens/index.ts +61 -1
- client/src/shared/api/GLTR_API.ts +55 -0
- client/src/shared/lang/translations.ts +5 -0
- experiments/tiny_nla/PLAN_v2.md +406 -0
- experiments/tiny_nla/batch_teacher_labels.py +105 -0
- experiments/tiny_nla/eval_final.py +179 -0
- experiments/tiny_nla/eval_roundtrip.py +157 -0
- experiments/tiny_nla/expand_dataset.py +82 -0
- experiments/tiny_nla/expand_dataset_v2.py +61 -0
- experiments/tiny_nla/gen_demo_html.py +43 -0
- experiments/tiny_nla/gen_opencode_only.py +99 -0
- experiments/tiny_nla/generate_data.py +397 -0
- experiments/tiny_nla/generate_teacher_labels.py +178 -0
- experiments/tiny_nla/generate_teacher_labels_opencode.py +207 -0
- experiments/tiny_nla/generate_teacher_labels_v2.py +186 -0
- experiments/tiny_nla/infer_tiny_nla.py +316 -0
- experiments/tiny_nla/nla_meta.yaml +48 -0
- experiments/tiny_nla/postprocess_dataset.py +112 -0
- experiments/tiny_nla/smoke_stage0.py +401 -0
- experiments/tiny_nla/train_ar.py +301 -0
- experiments/tiny_nla/train_av.py +331 -0
- experiments/tiny_nla/train_rl_grpo.py +404 -0
- experiments/tiny_nla/train_sft_v2.py +451 -0
- requirements.txt +1 -0
- server.py +1 -0
- server.yaml +61 -0
- 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> (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 |
+
```
|