使用大型语言模型在交互式机器学习笔记本中建议代码编辑
软件工程
2025-01-17 v1 计算与语言
机器学习
摘要
机器学习开发者频繁使用交互式计算笔记本(如 Jupyter notebooks)来承载用于数据处理和模型训练的代码。Jupyter notebooks 为编写机器学习流水线并交互式观察输出提供了便捷工具,然而由于笔记本的长度和复杂性,维护 Jupyter notebooks(例如添加新功能或修复错误)可能具有挑战性。此外,目前尚无与开发者在 Jupyter notebooks 上编辑相关的基准。为解决此问题,我们提出了首个包含 48,398 次 Jupyter notebook 编辑的数据集,该数据集源自 GitHub 上 792 个机器学习代码库的 20,095 次修订,并首次开展了使用 LLM 预测 Jupyter notebooks 中代码编辑的研究。我们的数据集捕获了单元格级和行级修改的细粒度细节,为理解机器学习工作流中的真实维护模式提供了基础。我们观察到,Jupyter notebooks 上的编辑高度局部化,代码库中的平均更改仅 166 行代码。尽管较大模型在代码编辑方面优于较小模型,但所有模型在我们的数据集上即使经过微调后准确率仍很低,这证明了真实世界机器学习维护任务的复杂性。我们的发现强调了上下文信息在提升模型性能方面的关键作用,并指明了提升大型语言模型在机器学习代码工程能力方面的有前景的方向。
引用
@article{arxiv.2501.09745,
title = {Suggesting Code Edits in Interactive Machine Learning Notebooks Using Large Language Models},
author = {Bihui Jin and Jiayue Wang and Pengyu Nie},
journal= {arXiv preprint arXiv:2501.09745},
year = {2025}
}