Mellum：编辑器中的快速小型代码补全模型

Mellum：编辑器中的快速小型代码补全模型

集成开发环境 (IDE) 的强大和精确的自动完成功能长期以来一直受到支持。然而，一个新的竞争者进入了这个领域：人工智能 (AI)。人工智能集成到 IDE 中正变得越来越普遍，这促使人们重新评估传统工具。虽然这些较旧的方法仍然有效，但它们可能不再提供开发人员现在期望的尖端帮助。

本文改编自 JetBrains 的一篇博客文章，深入探讨了 Mellum 的训练和功能，Mellum 是支持基于云的自动完成的模型。 JetBrains 以其为程序员量身定制的 IDE 和代码编辑器而闻名，它开发了 Mellum 以提供快速高效的代码自动完成功能。 Mellum 旨在在本地设备上有效运行，使其成为同类产品中最快、最紧凑的模型之一。一个概念上相似的项目是 Microsoft Phi。

应对显而易见的挑战

JetBrains 的工程师在追求 AI 驱动的自动完成功能时遇到了几个关键挑战：

• 速度和成本： 传统的聊天模型被证明是不切实际的，因为它们的计算成本很高且响应时间很慢。这些模型也缺乏对代码特定技术的了解，例如中间填充 (FIM) 或令牌修复。
• 输出格式： 旗舰聊天模型通常以不一致的格式生成数据，使得解析响应并将其无缝集成到编辑器中变得困难。
• 数据来源： 确定训练数据的来源并减轻潜在的版权侵权问题提出了一个重大障碍。

Mellum：概述

JetBrains 的开发团队意识到创建自己的模型是最佳方法。他们的目标是设计一种在质量、推理成本和延迟之间取得平衡的模型，同时在具有明确来源的数据上进行训练。初步研究表明，具有大约 40 亿个参数的模型可以为各种场景和用户提供有效的自动完成功能。此外，通过专门在代码上训练模型，他们可以建立一个专门的令牌词汇表，免受不相关数据的影响。

Mellum 模型的训练过程包括三个不同的阶段，每个阶段都贡献新的知识并提高生成的代码的质量。初始阶段涉及对大量单个文件进行基本预训练。第二阶段包括使用较小的一组专门示例来改进模型。人工智能反馈强化学习 (RLAIF) 用于第三阶段，以使模型适应 IDE 的特定特征并消除不需要的输出。

预训练

为了避免与模糊数据来源相关的潜在问题，该模型是从头开始训练的，需要全面介绍多种语言、编程语法、模式和核心概念。

数据集

预训练的主要数据来源是 TheStack。此数据集确保数据在法律上合理且在实践上有益。

预训练过程

在预训练期间，数据集被多次采样以达到大约 3 万亿个令牌。使用了 8192 个令牌的上下文窗口，数据集被分成这种大小的片段。中间填充 (FIM) 转换应用于每个片段中一半的文件，鼓励模型同时考虑前面的代码和后面的代码。这种技术非常类似于真实世界的代码生成场景。

预训练阶段是在由 16 个节点组成的集群上进行的，每个节点都配备了八个 H100 GPU。此阶段大约需要 15 天才能完成，从而产生了具有 40 亿个参数的 Mellum-base 模型。

预训练创建了一个通用的代码自动完成模型，具有对许多编程语言的广泛了解。但是，在此阶段，该模型仅经过训练来预测随机选择的文件段中的下一个令牌。如果没有其他上下文，该模型将缺乏对代码结构的了解，并且没有确定何时停止生成代码的机制。

微调阶段旨在解决这些限制。

上下文感知微调

增强的中间填充

与预训练（其中随机选择代码片段进行预测）不同，微调侧重于以更有意义的方式对代码进行分段，从而教导模型提取 ‘在野外’ 发生的代码片段。

专业示例

在实践中，代码自动完成需要了解周围的文件和更广泛的上下文，可能包括整个项目。

对于数据预处理，该公司启动了一个内部项目，代号为 Code Engine：一个跨平台 SDK 和一组控制台实用程序，用于直接从普通文件构建上下文，而无需完整的项目索引。此 SDK 部署在内部 MapReduce 集群上，用于处理数千个公共存储库，从而在合理的时间范围内生成了许多有用的训练示例。

找到正确的算法需要一些反复试验。

针对特定语言的调整

小型模型可以从针对特定语言的专业化中受益匪浅。虽然基础模型是在 80 多种语言上训练的，但大多数用户通常只使用一两种语言。为了解决这个问题，JetBrains 创建了多个专用模型：

• mellum-all： 支持 JetBrains IDE 中可用的大多数语言和方言，但自动完成质量低于专用模型。
• mellum-python： 专门研究 Python 和 Jupyter。
• mellum-kotlin： 专门研究 Java 和 Kotlin。
• mellum-web： 专门研究 Web 技术。

最后一步：RLAIF

最后，必须解决学习目标与用户期望不符的情况。一个额外的训练阶段，RLAIF — 人工智能反馈强化学习，用于解决此类问题。

该模型从用户交互中学习，并了解如何更好地反映用户偏好。

这种方法不仅提高了整体质量评分，还减少了烦人的生成工件的数量。

Mellum 有多好？

该模型对于其尺寸而言表现非常出色。以下是评估方式：

• 首先，该模型在内部基准测试中进行了评估，代号为 ‘JetBrains BigCode’。
• 然后在众所周知的公共基准测试（例如 SAFIM）上对其进行了测试。
• 最后，收集了功能的使用情况统计信息，并计算了用户指标。

离线评估

收集数据是一项复杂的任务，但是创建一种好的指标来比较原始建议和神经网络提出的新建议更具挑战性。我们进行了一项小型研究，最终确定了两个主要指标的组合：

EM：

• 完全匹配是一个非常受欢迎的想法。
• 如果完成的第一行与原始的第一行匹配，并且进行了最少的预处理，则认为预测是好的。

KK：

• 该指标以其作者的名字命名。
• 原始建议的行数除以建议完成的行数。

JetBrains BigCode

该模型是针对使用内部 JetBrains BigCode 工具获得的基准数据集进行评估的。

通过保持对数据集的完全控制，而不是依赖于公共基准，可以可靠地评估各种编码风格和实践的模型质量。

我们的 JetBrains BigCode 评估结果显示，质量与流行的模型相当，但 Mellum 更小、更高效。

单行建议的质量（EM 指标）

公共基准

该模型不仅在内部数据集上进行了评估，还在各种公共基准上进行了评估，例如多语言基准 SAFIM（语法感知中间填充）。

在线评估

主要指标称为已完成代码的比率 (RoCC)。它定义为使用代码自动完成编写的代码字符数与编辑器中代码总量的比率。

另一个重要指标是接受率 (AR)，它是计算为接受的建议数除以显示的所有建议数。

这是一段复杂的旅程，但 JetBrains 专家以尊严完成了它。最后，获得了一个通用模型和几个专用模型，这些模型通过 JetBrains AI 平台提供。它们现在在 JetBrains AI Assistant 中成功运行。

接下来是什么？

• JetBrains 的工程师目前正在开发用于 Web 开发语言的模型。它可能在不久的将来公开可用。
• 有计划同时增加参数数量和数据的多样性。编码中有许多不同的任务 — Mellum 也将能够执行这些任务。服务性能仍然是一个关键指标，因此模型的扩展将在合理的范围内。