Назад

Составление контекста для автодополнения всей строки кода

Антон Семенкин
anton.semenkin@jetbrains.com JetBrains Белград, Сербия

Ярослав Соколов
yaroslav.sokolov@jetbrains.com JetBrains Берлин, Германия

Евгения Ву
evgeniia.vu@jetbrains.com JetBrains Бремен, Германия

2024

Аннотация

Автодополнение кода --- одна из самых используемых функций интегрированной среды разработки (IDE), влияющая на повседневную жизнь разработчика программного обеспечения. Современные подходы к автодополнению кода перешли от композиции нескольких компонентов, основанных на статическом анализе, к конвейерам, включающим нейронные сети. Это изменение позволяет предлагать более длинные предложения кода (например, целую строку или даже фрагмент кода), сохраняя при этом относительно короткое время генерации. В JetBrains мы прилагаем много усилий для совершенствования рабочего процесса автодополнения, чтобы он был полезным и не отвлекающим для программиста. Нам удалось выпустить функцию Full Line Code Completion в IDE PyCharm Pro и доказать ее полезность в A/B-тестировании на сотнях реальных пользователей Python. В статье описывается наш подход к составлению контекста для модели Transformer, которая является ядром реализации функции. Кроме того, мы делимся нашими следующими шагами по улучшению функции и подчеркиваем важность нескольких исследовательских аспектов в этой области.

Ключевые слова: Автодополнение кода, Трансформеры, Составление контекста, Инжиниринг промптов, Интегрированная среда разработки, Программирование, Искусственный интеллект

1. Введение

Автодополнение кода --- одна из самых используемых функций IDE [1, 7], присутствующая в повседневной жизни разработчика ПО. Современные подходы к автодополнению кода перешли от композиции нескольких компонентов, основанных на статическом анализе, к конвейерам, включающим нейронные сети. Последние позволяют предлагать более длинные предложения кода (например, целую строку кода или даже фрагмент), сохраняя при этом относительно короткое время генерации. Хотя многострочные предложения являются мейнстримом в автодополнении на основе ИИ, есть свидетельства, что однострочные предложения могут быть более предпочтительными [6, 8] для разработчиков. Наше внутреннее исследование пользователей также показывает, что многострочное дополнение --- это ситуационно необходимая функция, однако однострочное дополнение удобно часто, что подчеркивает его важность.

В JetBrains мы прилагаем много усилий для совершенствования рабочего процесса автодополнения, чтобы он был полезным и не отвлекающим для программиста. В нашей команде Full Line Code Completion проект сосредоточен на генерации предложений кода, ограниченных одной строкой. Чтобы решить потенциальные проблемы конфиденциальности, связанные с передачей данных облачным провайдерам, мы решили работать исключительно на локальном устройстве конечного пользователя. Модель запускается нативно, что позволяет нам соблюдать строгие ограничения по времени выполнения, обеспечивая плавный пользовательский опыт. Однако это ограничение локальной обработки накладывает предел на размер модели нейронной сети, которую мы можем развернуть. Наши текущие оценки предполагают, что верхняя граница составляет около 1 миллиарда параметров, чтобы задержка и потребление памяти оставались приемлемыми. Следовательно, максимизация эффективности модели в рамках этого ограничения по размеру чрезвычайно важна. Таким образом, составление контекста, также известное как инжиниринг промптов, становится важнейшим аспектом реализации функции.

Недавно наша команда выпустила функцию Full Line Code Completion (FLCC) для различных IDE на основе IntelliJ, включая PyCharm Pro [5]. Функция оказалась значительно полезной в A/B-тестировании и работает безупречно, естественным образом улучшая рабочий процесс разработчика. Наш текущий подход позволяет интегрироваться в существующий рабочий процесс разработчика, не нарушая привычные паттерны.

В этой статье мы сосредоточимся на:

2. Full Line Code Completion в настоящий момент

2.1 Текущая настройка

В основе функции Full Line Code Completion мы используем авторегрессионные языковые модели, подобные GPT и LLaMA, которые предсказывают последовательность токенов на основе предыдущего контекста. Когда автодополнение кода вызывается (автоматически или действием пользователя), мы составляем контекст для модели следующим образом.

Поскольку наша цель --- работа на устройстве конечного пользователя, мы очень заботимся об эффективности составленного контекста. Работа на чьем-то локальном компьютере накладывает ограничения на работу на CPU с относительно коротким контекстом, поэтому использование оперативной памяти используемой модели Transformer не резко возрастает. Итак, среди других деталей реализации, которые мы придумали, мы более подробно описываем два важных решения.

Во-первых, мы выполняем обрезку пробелов, что означает не только удаление пустых строк, но и удаление ведущих пробелов в каждой строке. Для языков, чувствительных к отступам, таких как Python, ведущие пробелы составляют важную часть программы, поскольку указывают на изменение области видимости. Поэтому для таких языков мы заменяем каждое изменение области видимости соответствующим специальным токеном: <SCOPE_IN> и <SCOPE_OUT>. Таким образом, мы избегаем наличия токенов, которые отличаются только количеством пробелов в начале. Типичными примерами таких токенов были бы:

Во-вторых, мы не используем стандартное BPE для токенизации. Вместо этого мы используем «длинные токены», что по сути является модифицированным подходом BPE, который позволяет объединять токены через символы пробела, но запрещает объединение через конец строки. Этот подход позволяет нам создавать лучшие представления строк в исходном коде и сжимать его более эффективно. Например, следующие повторяющиеся «идиомы» в Python сжимаются до одного токена:

Объединяя все детали вместе, такой подход позволяет нам получить из контекста следующую информацию: конкретный диалект языка или тип файла, в котором вызывается автодополнение; общую информацию о модуле и особенностях файла; и компактное представление кода выше каретки для дополнения.

С описанным составлением контекста нам удалось создать функцию автодополнения кода в IDE, которая доказала свою полезность во время A/B-тестирования на сотнях реальных пользователей. Доля завершенного кода (от общего количества кода, набранного в редакторе) у пользователей с включенным Full Line Code Completion увеличилась в 1,5 раза по сравнению с пользователями, у которых эта функция не была включена. Кроме того, мы также получили явные отзывы от десятков пользователей. Отзывы были в основном положительными, вплоть до выражений, что функция предпочтительнее некоторых облачных инструментов на основе больших языковых моделей, но также описывались потенциальные точки роста. Пользователи явно хотят, чтобы функция работала лучше с учетом существующего контекста их работы, будь то точное местоположение в текущем файле, недавно открытые файлы или весь проект. Нашей команде удалось создать быструю, полезную и хорошо принятую функцию, которая уже внедрена в PyCharm Pro [5] и DataSpell [4].

2.2 Настройка оценки

Для последовательного улучшения Full Line Code Completion мы разработали различные механизмы обратной связи, охватывающие как офлайн, так и онлайн методы оценки. Офлайн-оценка проводится с использованием набора данных, содержащего исходный код --- метод, часто используемый в академической среде из-за типичного отсутствия доступа к широкой пользовательской базе. И наоборот, онлайн-оценка реализуется путем сбора данных от реальных пользователей, что делает ее предпочтительным подходом для разработки продукта. В FLCC мы сначала оцениваем производительность системы с помощью офлайн-оценки; однако наш акцент преимущественно лежит на онлайн-оценке, которую мы кратко описываем.

В JetBrains несколько раз в год проводится программа раннего доступа (EAP), позволяющая пользователям бесплатно загружать «экспериментальные» версии IDE. Во время EAP мы разделяем пользователей на несколько групп, предоставляем им разные версии некоторых функций и отслеживаем ряд метрик для каждой группы, также принимая во внимание статистическую значимость наблюдаемых результатов.

Такое A/B-тестирование было проведено нашей командой для функции автодополнения кода в PyCharm Pro. Осенью 2023 года мы изучали не только пользователей с FLCC и без нее, но и сравнивали разные реализации самой FLCC. Наша ключевая метрика --- доля завершенного кода (среди всего кода, написанного в редакторе) --- увеличилась в 1,5 раза у пользователей с включенным FLCC по сравнению с теми, у кого было только стандартное автодополнение. Мы также наблюдали, что пользователи не редактируют выбранный фрагмент кода сразу после его вставки в свой код. Наконец, мы не наблюдали никакого снижения производительности IDE при включении Full Line Code Completion.

2.3 Наши запланированные дальнейшие шаги

Помимо выпуска в других продуктах на основе IntelliJ, наша следующая цель --- решить проблему составления контекста, чтобы улучшить Full Line Code Completion еще больше.

Недавно мы начали экспериментировать с контекстами большего размера, чем раньше, изучая, как размер промпта связан с качеством предложений в нашей настройке с относительно небольшими нейронными сетями. Мы обнаружили, что можем увеличить контекст в несколько раз, не слишком беспокоя пользователей потреблением ресурсов машины. Для выпуска PyCharm Pro в декабре 2023 года использовалась модель с максимальным размером контекста 384 токена. Однако для предстоящих выпусков в 2024 году мы провели эксперименты для моделей с максимальным размером контекста 1536 токенов. Такое изменение стало возможным благодаря использованию LLaMA вместо GPT. Согласно нашей офлайн-оценке, нам удалось повысить качество автодополнения кода на 40%, сохранив задержку почти неизменной для большинства реальных сценариев. Мы продолжаем проводить онлайн-оценку реализованного изменения, чтобы оценить его влияние.

Более длинные контексты позволяют нам двигаться вперед к «умным» техникам составления контекста, которые включают методы fill-in-the-middle [3] и retrieval-augmented generation [2]. При большем размере контекста в промпт можно включить не только строки выше фактической позиции каретки, но и ниже нее. Кроме того, информацию можно извлекать из недавно открытых файлов, чтобы соответствующие фрагменты кода добавлялись в контекст модели. Мы уже провели несколько офлайн-экспериментов с многообещающими результатами (около 10% увеличения целевой метрики), поэтому следующий шаг --- доработать соответствующий опыт на стороне плагина и провести онлайн-эксперименты.

3. Исследовательские темы для Full Line Code Completion

3.1 Проведенные нами эксперименты

Существует известная по отзывам пользователей проблема «короткого контекста» --- модель быстро «забывает» содержимое файла выше текущей позиции набора. Эта проблема не ограничивается информацией в начале длинного файла, но также может распространяться на методы текущего класса, которые находятся значительно выше позиции каретки. Мы придумали подход, который может смягчить эту проблему.

Подход основан на следующей идее. Наиболее релевантный контекст для текущей позиции набора --- это код текущего метода или функции, объединенный с объявлениями других методов и функций, существующих в области видимости. Этот подход отражает рабочий процесс программистов: быстро получить представление о том, что реализовано в текущем классе, и обратить внимание на ближайший фрагмент кода. Для реализации мы составили контекст следующим образом:

Несмотря на следование интуиции рабочего процесса реального человека, подход не показал никаких положительных результатов в ходе экспериментов. Поэтому мы отказались от этого направления исследований.

3.2 Открытые проблемы

Мы твердо уверены, что наполнение контекста релевантным кодом --- это продуктивная точка роста для нейронных систем автодополнения кода (особенно тех, которые работают на устройстве конечного пользователя). Из исследования больших языковых моделей ясно, что инжиниринг промптов имеет значение, поэтому мы с нетерпением ждем обширных исследований, проведенных в этой области и для не очень больших языковых моделей. Область, наиболее интересная для нас, включает следующее.

Программисты часто переключаются между реализацией класса и его использованием, а также могут разрабатывать несколько вариантов использования одновременно, поэтому контекст недавно использованных файлов имеет значение. Ключевая проблема для решения этой задачи --- собрать релевантные данные для обучения моделей. Мы, конечно, не можем собирать такие данные от пользователей IntelliJ, кроме как при потенциальном явном согласии на участие в такой программе сбора данных. Данные потенциально можно собирать в анонимизированной форме, обучение таких моделей может проводиться федеративным образом, или существующие модели автодополнения кода можно дообучить на наборе данных ограниченного размера, собранном добровольцами. Мы ищем любые идеи или реализованные решения в этом направлении.

Еще одно перспективное направление для нейросетевых систем автодополнения кода --- это методы, подобные RETRO [2], которые интегрируют систему поиска в модель Transformer. Такие системы позволяют автоматически искать релевантные фрагменты кода среди фиксированной коллекции файлов исходного кода. Полученные фрагменты кода затем можно использовать для обогащения контекста модели Transformer, чтобы она генерировала лучшие предложения кода. Обширное исследование различных методов такого обогащения контекста представляет для нас существенный интерес.

Мы знаем, что исследователи часто не имеют доступа к наборам данных, включающим паттерны использования IDE пользователями. Например, объединение наборов данных смежных вкладок IDE представляет собой сложную задачу, связанную с проблемами конфиденциальности, когда данные собираются от реальных пользователей. Мы открыны для сотрудничества, которое может привести к созданию наборов данных описанного характера, которые могли бы быть полезны как исследователям, так и практикам, будучи полностью юридически корректными.

4. Заключение

Подводя итог, наша работа над Full Line Code Completion в JetBrains подчеркивает эффективное использование компактной, контекстно-зависимой модели, подобной GPT, адаптированной для эффективной работы на устройствах конечных пользователей. Этот подход показал многообещающие результаты в улучшении рабочих процессов написания кода в IDE JetBrains. Взглядывая вперед, мы с энтузиазмом относимся к изучению более информативных контекстов и ожидаем ценного вклада от широкого исследовательского сообщества в продвижении нейронных систем автодополнения кода.