引言

在现代机器学习（ML）开发流程中，交互式计算笔记本（如 Jupyter Notebook）因其灵活、直观的特点而被广泛采用。但随着项目规模和复杂度的提升，笔记本中 ML 管道代码的维护难度也显著增加：新增功能、修复缺陷，甚至是简单的重构，都可能涉及大量重复或分散的编辑工作。面对这一挑战，“如何自动化、智能化地对笔记本代码进行编辑”成为亟待解决的问题。

本文基于“Learning to Edit Interactive Machine Learning Notebooks”一文及其开源仓库内容，向专科及以上技术读者通俗介绍该项目的整体思路、数据集构建、实验方法以及关键发现。全文严格依托原始仓库说明，不另行补充外部知识，力求清晰、自然地呈现项目核心价值和可复现流程。在不追逐短期流量的前提下，重点呈现真实有效的技术细节与实用思考。

背景与动机

• 维护挑战
  传统的 ML 开发往往分为脚本化与模块化两条路径，而 Jupyter Notebook 将数据处理、模型训练、可视化展示等步骤串联于同一文档中，极大提升了原型设计和实验验证效率。然而，随着项目迭代，笔记本文件往往会积累大量旧版代码、注释甚至废弃单元，导致维护成本飙升。

• 缺乏基准
  尽管大型语言模型（LLM）在通用代码生成与补全方面表现优异，目前业界尚无针对笔记本编辑任务的基准数据集，也缺乏系统评估模型在真实笔记本维护场景下的能力指标。

基于以上动机，研究团队提出了首个面向笔记本编辑的基准数据集，并在此基础上探索 LLM 的微调与推理策略，以期为未来的智能化笔记本维护方案奠定基础。

数据集概览

研究者从 792 个与机器学习相关的 GitHub 仓库中提取了笔记本文件的历史修改记录，形成了包含 48,398 条笔记本编辑样本的数据集，涵盖 20,095 次独立提交（revisions）。其中，每条记录都包含文件级和单元格级（cell-level）的详细修改信息。

• 原始文件 commits.jsonl
  每行记录一个文件级别的编辑，包括：仓库名、提交信息、提交哈希、文件名、编辑前后内容等字段。
• 过滤后文件 commits-filtered.jsonl
  在原始数据基础上剔除了提交信息长度低于 3 字（英文）的记录，便于后续 LLM 学习。

数据收集流程

1. 仓库列表获取

使用 python_fetch.py 脚本调用 GitHub API，筛选出热门的 ML 仓库，并输出 top_1000_python_repos.json —— 一个按人气排序的前 1000 名 Python 仓库清单。

python python_fetch.py
# 输出：top_1000_python_repos.json

2. 历史提交克隆与提取

脚本 change_stat.py 会遍历 top_1000_python_repos.json 中的仓库，克隆仓库后逐个扫描提交历史，对每次提交中的笔记本文件（.ipynb）进行对比，提取编辑前后的内容并生成 commits.jsonl。

python change_stat.py
# 输出：commits.jsonl

• 字段说明

  • repo：仓库名称
  • commit：提交哈希
  • message：提交说明
  • file：文件路径
  • old / new：编辑前后文件内容

3. 数据清洗与划分

脚本 split.py 用于去除重复记录，并将样本按照比例划分为训练集、验证集和测试集，生成对应的索引文件 train_index.txt、val_index.txt、test_index.txt。

python split.py
# 输出：train_index.txt, val_index.txt, test_index.txt

处理完后，重新生成无重复记录的 commits.jsonl，确保数据质量和划分一致性。

代码编辑实验

微调模型（Finetuning）

为了让 LLM 更好地适应笔记本编辑任务，项目提供了基于 LoRA（Low-Rank Adaptation）的微调脚本 train_baseline.py。用户可以根据需求选择不同规模的基础模型（1.3B 或 6.7B 参数量）及不同训练轮次。

• 关键参数

  {
    "tokenizer.path": "deepseek-ai/deepseek-coder-6.7b-instruct",
    "model.path":     "deepseek-ai/deepseek-coder-6.7b-instruct",
    "model.rank":     8,
    "alpha":          16,
    "learning_rate":  3e-4,
    "training.epochs": 3,
    "per_device_train_batch_size": 8,
    "gradient_accumulation_steps": 2
  }
  

• 执行命令

  python train_baseline.py --config lora_config.json
  

• 输出结果
  训练完成后，在对应目录（如 1.3b_file/final）下得到微调后的模型权重。

基线推理（Inference）

基线脚本 baseline.py 支持多种推理模式，包括：

1. 零样本（Zero-shot）
   不提供示例，仅根据提交消息和原始代码预测编辑结果。
2. 一样本（One-shot）
   在提示中加入单个示例，提高模型对编辑方式的理解。
3. 五样本（Five-shot）
   提供五个示例，进一步增强提示信息。

示例调用（文件级零样本）：

from baseline import generate_code_change_whole_file

for d in test_data:
    tgt = generate_code_change_whole_file(
        tokenizer, model,
        commit_message=d.message,
        original_code=d.old,
        cell_diff=d.cell_diff
    )
    # 保存 tgt 到指定路径

用户可通过修改脚本中的 ftt 参数（表示第几次微调）和 file 或 cell 模式，灵活切换不同模型与推理场景。

数据集统计

为了量化数据集覆盖度，项目提供以下统计脚本：

• 模型统计：model_stat.py
  统计不同模型（基础 vs. 微调）在不同任务上的 token 分布、样本长度等信息。
• 数据量统计：dataset_size_stat.py
  输出完整数据集及训练/验证/测试集的样本量统计。

执行示例：

python model_stat.py
python dataset_size_stat.py

通过可视化图表或表格，团队可以直观了解数据集特征，并为模型设计与训练提供依据。

评价指标

该项目选用多种指标评估模型在笔记本编辑任务中的表现，主要包括：

• BLEU / CodeBLEU
  比较生成代码与参考代码在 n-gram 级别的重合度。
• EditSim
  度量预测编辑与实际编辑在操作层面的相似度。
• ROUGE-L
  基于最长公共子序列计算文本相似度，适用于衡量整体内容一致性。

执行评估脚本：

python accuracy.py \
  --output_folder model/results/whole_file_zero_shot_1.3b \
  --expected_folder model/results/expected_whole_file
python finetune_score.py --tpe whole

最终会生成 .pkl 格式的评估结果，便于后续分析与对比。

实践要点与思考

1. 上下文信息的重要性
   实验表明，笔记本编辑高度依赖上下文：文件级的全局信息与单元格级的局部信息，均需被模型充分理解，才能做出准确的修改预测。
2. 数据质量 vs. 数量
   虽然数据量巨大，但粗糙的或噪声过多的提交会影响模型学习效果。过滤短提交、去重等数据清洗工作不可或缺。
3. 示例数的平衡
   提供更多示例（few-shot）固然能提升提示效果，但也带来更高的上下文窗口开销。选择恰当的示例数量，是实用部署时需重点考虑的问题。

总结

本文基于“Learning to Edit Interactive Machine Learning Notebooks”项目仓库，深入介绍了笔记本编辑场景下的数据集构建、微调与推理方法，以及关键统计与评估流程。项目填补了交互式笔记本维护领域的基准空白，为后续研究和应用提供了宝贵资源。读者可以根据本文所述步骤，复现数据收集与模型实验，或以此为基础，探索更高效、更准确的自动化编辑方案。

如需深入了解或获取数据集，请访问Zenodo 数据集链接 或查阅原始论文。