Jupyter Agents:訓練 LLM 以筆記本推理
Hugging Face 開發了一條流水線,用於訓練小型語言模型,使其能作為能在 Jupyter Notebook 中執行程式碼的資料科學代理。透過結合從 Kaggle Notebook 衍生的精選合成資料集與簡化的代理腳手架,團隊成功提升了 Qwen3-4B 模型在資料分析基準測試上的表現。
DABStep 基準測試:資料科學代理
為了衡量代理式資料科學能力的進展,Hugging Face 使用 DABStep 基準測試,該測試透過提供資料集並要求模型回答非平凡的資料問題,來評估代理在真實任務上的表現。此基準對目前的 LLM 來說相當具挑戰性;例如,Claude 4 Sonnet 在基準的困難任務上正確率不足 20%。
最佳化代理腳手架
代理的效能高度依賴腳手架——指導模型行為的程式碼與提示結構。Hugging Face 重新構築了 Jupyter Agent 的腳手架,使其輕量且符合模型需求:
- 簡化架構: 腳手架被縮減至約 200 行程式碼,且沒有外部依賴。
- 工具整合: 代理在 while 迴圈中運作,使用兩個主要工具:
code execution與final_answer。 - 效能提升: 此最佳化單獨就將 DABStep 簡易分割的正確率從 44.4% 提升至 59.7%。
Jupyter Agent 資料集流水線
為了提升小型模型,Hugging Face 使用 Datatrove 建立了多階段流水線,將約 2TB 的原始 Kaggle Notebook 轉換為高品質的 51k 合成 Notebook 訓練集(約 0.2B 代碼)。
1. 去重與資料取得
原始 Kaggle Notebook 進行去重,將容量從 2TB 縮減至 250GB。團隊接著使用 kagglehub 套件下載約 5TB 的相關資料集,過濾掉模型檢查點、多模態語料庫,以及超過 10GB 的資料集,以確保能放入 E2B 沙盒中。
2. 教育評分與過濾
使用 Qwen3-32B,團隊實作了「教育評分」系統,根據清晰度與教育價值將 Notebook 評分為 1–5,刪除約 70% 的 Notebook。再透過 LLM 為基礎的過濾移除額外 20%,以排除與資料分析無關或未使用資料集的 Notebook。
3. 合成 QA 與追蹤生成
為了建立可驗證的訓練資料,流水線產生以真實 Notebook 追蹤為基礎的合成問答 (QA) 配對。
- QA 生成: Qwen3-32B 產生自然的問題與答案,之後由第二個 LLM 驗證以防止幻覺。
- 追蹤生成: 使用 Qwen-3-Coder-480B-A35B-Instruct 產生乾淨、逐步的程式碼執行追蹤,以回答合成問題。
- 處理缺失資料: 當資料集不可用時,模型被提示充當有狀態的 Python 程式碼解譯器以模擬執行。
- 推理模擬: 為確保模型在程式碼單元之間提供評論,
comment欄位被設定為程式碼執行工具呼叫的必填項目。
訓練與結果
微調在 Qwen3-4B 上使用 TRL 函式庫進行。團隊發現全參數微調優於 PEFT(參數高效微調),且使用 assistant_loss_only=True 能提升效能。
效能提升
在合成資料集上訓練顯著提升了模型在 DABStep 簡易分割上的效能:
| 模型 | 訓練輪次 | DABStep(簡易) |
|---|---|---|
| Qwen-3-4B-Instruct-2507 (基礎) | 0 | 38.67% |
| Qwen-3-4B-Instruct-2507 (我們的腳手架) | 0 | 52.78% |
| Qwen-3-4B-Instruct-2507 | 2 | 63.89% |
| Qwen-3-4B-Instruct-2507 | 3 | 73.61% |
| Qwen-3-4B-Instruct-2507 | 5 | 75% |
| Qwen-3-4B-Instruct-2507 | 7 | 70.83% |
最終微調的 Qwen-4B 模型在簡易分數上較基礎模型提升了最高 36%,較腳手架基礎模型提升了 22%,同時在困難分數上也有改善。
開源釋出
Hugging Face 已向社群釋出以下資源:
- Jupyter Agent Dataset: 51k 合成 Notebook 的集合。
- 已微調模型:
jupyter-agent-qwen3-4b-instruct與jupyter-agent-qwen3-4b-thinking。
未來方向
團隊找出多條進一步改進的路徑,包括產生更具挑戰性的多步問題、擴大精選追蹤的規模、研究知識蒸餾,以及利用現有可驗證 QA 設定構建強化學習 (RL) 環境。