В мире интегрированных сред разработки (IDE) функции автозавершения давно играют важную роль, повышая продуктивность и точность написания кода. Однако с появлением искусственного интеллекта (ИИ) в этой области произошла настоящая революция. Интеграция ИИ в IDE становится все более распространенной, что заставляет переосмыслить традиционные инструменты. Хотя старые методы по-прежнему работают, они могут больше не обеспечивать ту передовую помощь, которую разработчики ожидают сегодня.
Эта статья, адаптированная из блога JetBrains, посвящена обучению и возможностям Mellum, модели, лежащей в основе облачного автозавершения. JetBrains, известная своими IDE и редакторами кода, разработанными специально для программистов, разработала Mellum для обеспечения быстрого и эффективного автозавершения кода. Mellum спроектирована для эффективной работы на локальном устройстве, что отличает ее как одну из самых быстрых и компактных моделей в своем классе. Концептуально схожим проектом является Microsoft Phi.
Решение очевидных проблем
Инженеры JetBrains столкнулись с рядом серьезных проблем в своем стремлении к автозавершению на основе ИИ:
- Скорость и стоимость: Традиционные чат-модели оказались непрактичными из-за высоких вычислительных затрат и медленного времени отклика. Эти модели также не учитывали специфические для кода методы, такие как заполнение середины (FIM) или восстановление токенов.
- Форматирование вывода: Флагманские чат-модели часто генерировали данные в несогласованных форматах, что затрудняло анализ ответов и их беспрепятственную интеграцию в редактор.
- Происхождение данных: Определение источника данных для обучения и смягчение потенциальных проблем с нарушением авторских прав представляло собой серьезное препятствие.
Mellum: Обзор
Команда разработчиков в JetBrains пришла к выводу, что создание собственной модели является оптимальным подходом. Их целью было разработать модель, которая бы сбалансировала качество, стоимость инференса и задержку, при этом обучаясь на данных с четким происхождением. Предварительные исследования показали, что модель с примерно 4 миллиардами параметров может обеспечить компетентные возможности автозавершения для широкого круга сценариев и пользователей. Кроме того, путем обучения модели исключительно на коде они могли создать специализированный словарь токенов, свободный от нерелевантных данных.
Процесс обучения модели Mellum состоит из трех отдельных этапов, каждый из которых вносит новые знания и улучшает качество генерируемого кода. Начальный этап включает в себя базовое предварительное обучение на большом корпусе отдельных файлов. Второй этап состоит в уточнении модели с помощью меньшего набора специализированных примеров. Обучение с подкреплением с обратной связью от ИИ (RLAIF) используется на третьем этапе для адаптации модели к специфическим характеристикам IDE и устранения нежелательного вывода.
Предварительное обучение
Чтобы избежать потенциальных проблем, связанных с неоднозначным происхождением данных, модель была обучена с нуля, что потребовало всестороннего знакомства с многочисленными языками, синтаксисом программирования, паттернами и основными концепциями.
Набор данных
Основным источником данных для предварительного обучения был TheStack. Этот набор данных гарантирует, что данные являются как юридически обоснованными, так и практически полезными.
Процесс предварительного обучения
Во время предварительного обучения набор данных был выбран несколько раз, чтобы достичь примерно 3 триллионов токенов. Использовалось контекстное окно размером 8192 токена, при этом набор данных был разделен на фрагменты этого размера. Преобразование заполнения середины (FIM) было применено к половине файлов в каждом фрагменте, побуждая модель учитывать как предшествующий, так и последующий код. Этот метод тесно имитирует реальные сценарии генерации кода.
Фаза предварительного обучения проводилась на кластере из шестнадцати узлов, каждый из которых был оснащен восемью графическими процессорами H100. Этот этап занял примерно 15 дней, в результате чего была получена базовая модель Mellum с 4 миллиардами параметров.
Предварительное обучение создает универсальную модель автозавершения кода с обширными знаниями многих языков программирования. Однако на этом этапе модель обучена только предсказывать следующий токен в случайно выбранном сегменте файла. Без дополнительного контекста модель не осведомлена о структуре кода и не имеет механизма для определения, когда следует прекратить генерацию кода.
Этап тонкой настройки предназначен для устранения этих ограничений.
Тонкая настройка с учетом контекста
Улучшенное заполнение середины
В отличие от предварительного обучения, где фрагменты кода выбираются случайным образом для прогнозирования, тонкая настройка концентрируется на сегментировании кода более значимым образом, обучая модель извлекать фрагменты кода, которые встречаются ‘в дикой природе’.
Специализированные примеры
На практике автозавершение кода требует понимания окружающих файлов и более широких контекстов, возможно, охватывающих целые проекты.
Для предварительной обработки данных компания запустила внутренний проект под кодовым названием Code Engine: кроссплатформенный SDK и набор консольных утилит, разработанных для создания контекста непосредственно из обычных файлов без необходимости полной индексации проекта. Этот SDK был развернут во внутреннем кластере MapReduce и использовался для обработки тысяч общедоступных репозиториев, генерируя множество полезных примеров для обучения в разумные сроки.
Поиск правильных алгоритмов потребовал некоторых проб и ошибок.
Настройка для конкретных языков
Небольшие модели могут получить большую выгоду от специализации для конкретных языков. Хотя базовая модель обучена на более чем 80 языках, большинство пользователей обычно работают только с одним или двумя. Чтобы решить эту проблему, JetBrains создала несколько специализированных моделей:
- mellum-all: Поддерживает большинство языков и диалектов, доступных в IDE JetBrains, но качество автозавершения ниже, чем у специализированных моделей.
- mellum-python: Специализируется на Python и Jupyter.
- mellum-kotlin: Специализируется на Java и Kotlin.
- mellum-web: Специализируется на веб-технологиях.
Заключительный этап: RLAIF
Наконец, необходимо решить случаи, когда цели обучения не соответствуют ожиданиям пользователей. Для решения таких проблем используется дополнительный этап обучения — RLAIF (Reinforcement Learning with AI Feedback) — обучение с подкреплением с обратной связью от ИИ.
Модель учится на взаимодействии с пользователем и понимает, как лучше отражать предпочтения пользователя.
Этот подход не только улучшает общий показатель качества, но и уменьшает количество раздражающих артефактов генерации.
Насколько хорош Mellum?
Модель исключительно хорошо работает для своего размера. Вот как она оценивалась:
- Во-первых, модель была оценена на внутреннем эталоне под кодовым названием ‘JetBrains BigCode’.
- Затем она была протестирована на известных общедоступных эталонах, таких как SAFIM.
- Наконец, была собрана статистика использования функций и рассчитаны пользовательские метрики.
Оффлайн оценка
Сбор данных — сложная задача, но создание хорошей метрики, которая сравнивает исходное предложение с новым, предложенным нейронной сетью, еще сложнее. Мы провели небольшое исследование и в конечном итоге остановились на комбинации двух основных показателей:
EM:
- Exact Match (точное совпадение) — очень популярная идея.
- Прогноз считается хорошим, если первая строка завершения совпадает с первой строкой оригинала, с минимальной предварительной обработкой.
KK:
- Метрика названа в честь ее авторов.
- Количество предложенных строк из оригинала, деленное на количество строк в предложенном завершении.
JetBrains BigCode
Модель была оценена по эталонному набору данных, полученному с помощью внутреннего инструмента JetBrains BigCode.
Поддерживая полный контроль над нашим набором данных, а не полагаясь на общедоступные эталоны, становится возможным надежно оценивать качество модели для различных стилей и практик кодирования.
Результаты нашей оценки JetBrains BigCode показывают качество на уровне популярных моделей, но Mellum меньше и эффективнее.
Качество однострочных предложений (метрика EM)
Общедоступные эталоны
Модель была оценена не только на внутреннем наборе данных, но и на различных общедоступных эталонах, таких как многоязычный эталон SAFIM (syntax-aware fill in the middle — заполнение середины с учетом синтаксиса).
Онлайн оценка
Основная метрика называется коэффициентом завершенного кода (RoCC). Он определяется как отношение символов кода, написанных с использованием автозавершения кода, к общему объему кода в редакторе.
Другой важной метрикой является коэффициент принятия (AR), который рассчитывается как количество принятых предложений, деленное на количество всех показанных предложений.
Это было сложное путешествие, но специалисты JetBrains прошли его с достоинством. В конце концов, была получена одна общая и несколько специализированных моделей, которые доступны через платформу JetBrains AI. Сейчас они успешно работают в JetBrains AI Assistant.
Что дальше?
- Инженеры JetBrains в настоящее время работают над моделью для языков веб-разработки. В ближайшем будущем она может стать общедоступной.
- Есть планы одновременно увеличить как количество параметров, так и разнообразие данных. В кодировании есть много разных задач — Mellum сможет выполнять и их. Производительность сервиса по-прежнему является ключевой метрикой, поэтому расширение модели будет в разумных пределах.