Человек купил в магазине ящик елочных игрушек. На следующий день приносит его обратно:
- Ваш товар бракованный.
- Что, игрушки битые?
- Нет, целые.
- Не блестят?
- Блестят.
- Так в чем же дело?
- НЕ РАДУЮТ!
Автодополнение всегда было силой и гордостью современных IDE. Сравнивая IDEA, VSCode и Eclipse, мы смотрим — оно вообще адекватные вещи пишет в выпадающей менюшке? Знает про Spring, про Next.js, про свежие апдейты в Питоне?
Но в последнее время, детерминированные алгоритмы больше не являются предметом хайпа. В мире IDE появился новый царь горы — AI, который обязан быть везде. Старые игрушки продолжают работать, но уже не радуют.
Этот текст — адаптированный перевод статьи из блога JetBrains (я постарался сделать текст более читабельным и менее похожим на политкорректный пресс-релиз). Мы посмотрим, как обучилась модель Mellum, лежащая в основе облачного автодополнения, и что из этого вышло.
Напоминаю, основной продукт JetBrains — это их IDE, редактор кода для программистов. Модель Mellum занимается автодополнением кода в редакторе, и является одной из самых быстрых и миниатюрных в своем классе. Она способна запускаться на устройстве (не только в облаке у арабских шейхов). Чем-то похожим по смыслу можно считать Microsoft Phi.
Исследование: что насчет очевидных вариантов?
Все началось с тестирования всем известных закрытых LLM через API, и их альтернатив с открытым исходным кодом. Инженеры JetBrains уперлись в ряд проблем:
Для привычной моментальной скорости срабатывания автодополнения, типичные чат-модели оказались непрактичными: очень дорогими и тормозными. Они ничего не знали про типичные для кода приемы типа fill-in-the-middle (FIM) или token healing. Ну как "привычной" — все знают, что IDEA тормозит, но не настолько же.
Флагманские чатовые модели имеют тенденцию вываливать данные в непонятном формате, что затрудняет правильную обработку ответа и вставку текста в редактор.
Одной из самых больших проблем — непонятно, откуда берутся исходные данные, и как не получить иск в суд от изначального владельца. Кое-кто из вендоров предлагает гарантии, но JB решили с этим не заморачиваться.
Mellum: обзор с высоты птичьего полета
Тут разработчики дошли до мысли, что правильный путь — делать свою модельку. Нужно обучить модель с разумным балансом качества, затрат на вывод и задержкой. И обучить ее на данных понятного происхождения. Предварительное исследование показало, что даже с 4 миллиардами параметров, вполне можно получить приличное дополнение для большинства сценариев и пользователей. Кроме того, если модель изначально учить только на коде, можно получить специфический словарь токенов, не замусоренный другими способами применения.
Для обучения модели придумали трехэтапный процесс, где каждый этап привносит новые знания и улучшает качество генерации. Все началось с базового предварительного обучения на большом корпусе отдельных файлов, а результат заполировали небольшим количеством специализированных примеров. Для адаптации к специфике IDE и очистке от мусорных результатов, использовался reinforcement learning with AI feedback (RLAIF).
Shameless plug. Хотите свежего мяса и пейперов про нейронки? Телега @1red2black.
Предварительное обучение
Чтобы избежать рисков, связанных с использованием непонятных данных, обучение делалось с нуля. В смысле, вообще с самого нуля. Нужно было познакомить модель с кучей языков, научить синтаксису, шаблонам и концепциям программирования в целом.
Набор данных
Основным источником данных стал TheStack. Он не то чтобы актуален, но проблему можно решить простым лайфхаком: собрать свежий код из публичных репозиториев, оставляя только тот, где есть хорошая лицензия и нет персональных данных, за которые в Европе можно получить по башке. Это гарантирует, что данные не только нравятся юристам, но еще и полезны на практике.
Процесс предварительного обучения
Для предварительного обучения, набор данных выборочно использовался несколько раз, с целью получить приблизительно 3 триллиона токенов. Использовалось контекстное окно размером 8192 токена (набор данных распилили на фрагменты соответствующего размера). Для половины файлов в каждом фрагменте применили трансформацию fill-in-the-middle (FIM). Файлы делились на три части: префикс, середина и суффикс. Затем эти части переставлялись так, чтобы модель научилась предсказывать отсутствующий средний сегмент на основе окружающего контекста. Эта техника побуждает модель учитывать как предшествующий, так и последующий код, что лучше имитирует реальное использование при кодогенерации.
Предварительное обучение проводилось на кластере из шестнадцати узлов. На каждом узле — восемь GPU H100. Для завершения этапа потребовалось около 15 дней. Результатом стала 4-миллиардная модель Mellum-base. Для сравнения, 100-миллионные модели, которые локально работают в IDE, обычно обучаются около недели, всего на одном узле с восемью штуками H100.
Предварительное обучение создает модель автодополнения кода общего назначения с обширными знаниями о многих языках программирования. Однако, на данном этапе модель обучена достижению лишь одной простой цели: предсказанию следующего токена в случайно выбранном сегменте файла. Без дополнительного контекста модель не осознает структуру кода и не понимает, когда прекратить генерацию.
Эти ограничения призван устранить этап файн-тюнинга.
Контекстно-зависимый тюнинг
Улучшенный fill-in-the-middle
В отличие от предварительного обучения – где фрагменты кода для предсказания выбираются случайно – файнтюн фокусируется на нарезке кода более осмысленным способом. То есть, нужно научить ее извлекать фрагменты кода, которые встречаются "в дикой природе".
Посмотрите на изображение. Код, показанный синим цветом – то, что мы просим модель предсказать. Во втором примере справа, выбранный фрагмент укладываются внутри одной функции. Это более похоже на правду. Это более качественный пример для обучения.
Специализированные примеры
Даже с улучшенным FIM, мы всё ещё остаемся в рамках одного файла. Сложно найти живого разработчика, который хочет всё писать в одном файле. На практике, автодополнение кода требует понимания окружающих файлов и более широкого контекста — возможно, контекста всего проекта.
Для предварительной обработки данных, в компании запустили внутренний проект под кодовым именем Code Engine: кроссплатформенный SDK и консольные утилиты, разработанные для построения контекста непосредственно из обычных файлов, без необходимости полноценной индексации проекта. Этот SDK запустили на внутреннем MapReduce-кластере, прогнали через него тысячи публичных репозиториев, и за достаточно разумное время погенерили кучу полезных для обучения примеров.
Поиск правильных алгоритмов потребовал некоторых проб и ошибок. Какие были проблемы?
Сортировка файлов, которые построчно наиболее близки по индексу Жаккара
Включение файлов из операторов импорта
Создание полной карты репозитория
…и многое другое.
Тюнинг для конкретных языков
Идея в том, что маленькие модели могут значительно выиграть от специализации под конкретные языки. Хотя базовая модель уже обучена на более чем 80 языках, большинство пользователей обычно работают только с одним или двумя (например, Java и Python). Настройка модели под конкретный язык позволяет лучше улавливать особенности конкретного языка, библиотеки, или даже подходы к разработке:
mellum-all – поддерживает большинство языков и диалектов, доступных в IDE от JetBrains, но качество автодополнения хуже, чем у специализированных моделей;
mellum-python – специализируется на Python и Jupyter;
mellum-jotlin – специализируется на Java и Kotlin;
mellum-web – специализируется на веб-технологиях.
Последний шаг: RLAIF
И наконец, нужно подчистить те случаи, когда цели обучения не совпадают с ожиданиями пользователей (иначе говоря, пользователь считает, что модель сгенерировала дичь). Например, с при обучении вполне допустимо генерировать филлеры вроде TODO("Not implemented") — в публичных репозиториях это встречается повсюду. Но пользователь вряд ли будет счастлив, если такие филлеры будут в автодополнении.
Для решения таких проблем применяется дополнительный этап обучения, RLAIF — бучение с подкреплением на основе обратной связи от ИИ. В него включаются синтетические данные, созданные на основе правил и настроек, сгенерированных моделью. Полученный датасет состоит из двух больших категорий:
Пары дополнений на основе правил: берутся реальные примеры из предыдущего набора данных, но намеренно ухудшаются. Осмысленный код заменяется популярными заполнителями типа pass, TODO и комментариями .
Пары дополнений, сгенерированные Mellum: генерируются два дополнения – одно с низкой температурой (более детерминированное и качественное) и другое с высокой температурой (более непредсказуемое или более низкого качества). Для ранжирования этих двух дополнений используется внешняя LLM. На выходе — пара размеченная маркерами positive и negative.
Модель учится на них и понимает, как лучше отражать предпочтения пользователей. Сейчас используется алгоритм direct preference optimization (DPO), который делает модель более склонной к генерации позитивных примеров.
Такой подход не только повышает общую оценку качества, но и уменьшает число раздражающих артефактов генерации.
Насколько хорош Mellum?
Спойлер: модель работает невероятно хорошо для своего размера! Вот как ее оценивали:
Сначала модель оценили на внутреннем бенчмарке под кодовым именем "JetBrains BigCode".
Затем ее проверили на известных публичных бенчмарках, таких как SAFIM.
Наконец, подняли статистику использования фичей и посчитали пользовательские метрики.
Ниже циклы обратной связи разбиты на две категории: офлайн-оценка (с предопределенными наборами данных) и онлайн-оценка (с реальными пользовательскими данными).
Офлайн-оценка
Собрать данные — сложная задача, но еще сложнее создать хорошую метрику, которая сравнивает изначальное дополнение с новым, предложенным нейронкой. Мы провели небольшое исследование и в итоге остановились на комбинации двух основных метрик:
EM:
Точное совпадение (Exact match) — очень популярная идея.
Предсказание считается хорошим, если первая строка дополнения совпадает с первой строкой оригинала, с минимальной предварительной обработкой.
KK:
Метрика была названа в честь её авторов, Karol и Katya, которые не только придумали ее, но и вручную аннотировали пороговый датасет.
Число предложенных строк из оригинала, деленное на число строк в предложенном дополнении.
Используются и другие методы оценки: например, chrF и LLM-as-a-Judge, для сравнения вариаций. Мы пока не будем про них писать, но в будущем — все возможно.
JetBrains BigCode
Модель оценивалась по эталонному набору данных, полученному с помощью внутреннего инструмента JetBrains BigCode. Набор охватывает наиболее популярные языки, поддерживаемые Mellum, среди них — Python, Kotlin и Java. В нем удалили пересечения с исходным обучающим набором данных, и преобразовали в FIM-структуру с контекстами, собранными с помощью Code Engine.
Датасет JetBrains BigCode крут способностью разделять фрагменты кода на основе различных характеристик: тема репозитория, популярность, возраст, является ли код основным или тестовым, когда его в последний раз меняли... Во-первых, хорошая модель должна демонстрировать сильные результаты по всем этим срезам сразу, а не только по их небольшому подмножеству. Во-вторых, в публичных бенчмарках могут затесаться обучающие данные, и такая оценка получится плохой и грязной. Чтобы смягчить риск, срезы берутся по возрасту репозитория и активности в нем — это лучше отражает реальную жизнь.
Сохраняя полный контроль над нашим набором данных, а не полагаясь на публичные бенчмарки, становится возможным надежно оценивать качество модели для различных стилей и практик кодирования. Экспериментам с JetBrains BigCode показывают, что сложная природа базовой модели дает более качественный результат. Вначале стоит получить сильную базовую модель, потом подтюнить её, и только под конец заалайнить с помощью DPO.
Результаты нашей оценки JetBrains BigCode показывают качество на уровне с популярными моделями, но при этом Mellum меньше и эффективнее. Конечно, модели большего размера превосходят Mellum, но и затраты затраты на их обслуживание слишком велики.
Публичные бенчмарки
Модель оценивали не только на внутреннем наборе данных, но и на различных публичных бенчмарках, таких как мультиязычный бенчмарк SAFIM (syntax-aware fill in the middle). В этой статье данных нет, но в будущем они могут появиться.
Важно помнить, что бенчмарки, такие как JetBrains BigCode и SAFIM, несмотря на их полезность для массовой офлайн-оценки, далеко не полностью отражают опыт реального разработчика, использующего модель. А ведь, в конечном счете, именно для них всё это и делается.
Онлайн-оценка
Из статистики вычленяется несколько полезных пользовательских метрик. Используется тот же конвейер, который описан в статье "Full Line Code Completion: Bringing AI to Desktop".
Основная метрика называется коэффициентом завершенного кода (ratio of completed code, RoCC). Она определяется как отношение символов кода, написанных с использованием автодополнения кода, к общему количеству кода в редакторе. Основная идея и мотивация просты: чем больше больше кода в редакторе пользователь написал через автодополнение, тем лучше к нему пользователь относится. При этом, точно известно, откуда взялся код: это позволяет, например, отличить общий RoCC в редакторе, и специфичный только для генерации из Mellum.
Еще одной важной метрикой является частота принятий (acceptance rate, AR), которая вычисляется как количество принятых предложений, деленное на количество всех показанных предложений. Эта метрика широко используется в сообществе и довольно интуитивная: чем чаще пользователи соглашаются с предложениями, тем лучше.
Вот немного онлайн-метрик для различных популярных языков:
RoCC (Mellum + стандартное дополнение) | RoCC (только Mellum) | AR | |
|---|---|---|---|
Java | 46% | 23% | 35% |
Kotlin | 45% | 25% | 31% |
Python | 32% | 23% | 35% |
JS/TS | 39% | 23% | 32% |
C# | 45% | 18% | 32% |
Go | 45% | 30% | 44% |
PHP | 40% | 26% | 34% |
Rust | 37% | 24% | 35% |
Выводы
Это был сложный путь, но специалисты JetBrains его прошли с достоинством. В конце получилась одна общая, и несколько специализированных моделей, которые доступны через платформу JetBrains AI. Сейчас они успешно работают в JetBrains AI Assistant.
Что будет дальше?
Сейчас инженеры JetBrains работают над моделью для языков веб-разработки. Возможно, в ближайшем будущем она станет общедоступной.
Планируется одновременно увеличить и количество параметров, и разнообразие данных. В кодировании есть много разных задач — Mellum сможет выполнять и их тоже. Производительность сервиса все еще остается ключевой метрикой, поэтому расширение модели будет в разумных пределах.
Попробуйте кодировать с Mellum и поделитесь своими впечатлениями!
