CodeBert для автогенерации комментариев к коду

Код программ отличается от естественного языка из-за его формализма и строгости, однако, ничто не мешает воспринимать его как последовательность токенов и работать с ним, как с обычным языком. Существуют исследования, которые показали, что модель BERT, обученная на большом наборе данных, неплохо справляется с некоторыми задачами, связанными с обработкой программного кода.

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

Данные представлены в виде набора пар [функция — комментарий] для различных языков программирования (awesome Code Search Net Challenge dataset). Кстати говоря, этот набор изначально был создан не для этой задачи, однако его можно легко перепрофилировать под свои нужды.

Я не стану очищать данные, это описано здесь. Я же буду использовать уже предварительно обработанные данные в объеме 1 % от общего количества образцов в наборе, так как обучение модели занимает довольно много времени. Но, как можно будет убедиться в будущем, генерация комментариев даже на 1 % данных выглядит неплохо. Если у вас есть время и ресурсы, можете обучить модель на всём наборе и получить результаты получше.

Предварительно обученная модель, которую я буду использовать, взята из статьи исследовательского подразделения Microsoft. В этой модели также использовался набор данных CodeSearchNet, но вместо генерирования комментариев он использовался для обучения модели на основе RoBERTa удобному для восприятия представлению кода и естественного языка. Использование больших языковых моделей для представления текста удобным способом в настоящее время является обычной практикой, поскольку они показали свою эффективность для решения других задач.

Здесь я прописываю пути до файлов с данными и оборачиваю их в структуру для более удобного дальнейшего использования:

Инициализирую две вспомогательные функции: токенизации текста и записи DataFrame в JSON-файл, так как именно в таком формате требуется подавать данные для модели.

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

Далее создаю конфигурационный класс для модели и на его основе прописываю всю конфигурацию:

Теперь, когда данные обработаны и представлены в удобном формате, можно приступать к обучению. Сделаю этоОбучу модель на обучающей выборке. В качестве метрики использую BLEU-4 (четвёрка означает, что количество словесных n-gram = 4), которая распределена от 0 до 1, но в нашем примере будет использоваться BLEU-4 * 100%. Эта метрика используется в задачах машинного перевода, но и для генерации текста она также хорошо подходит. Если брать задачи машинного перевода, то даже для человека bleu = [0.6:0.7] — отличный результат, потому что каждый человек может перевести текст по-разному. Точности в единицу достигнуть почти невозможно.

Если посмотреть на исходную задачу, то, во-первых, модель должна сгенерировать текст, а во-вторых, это не просто текст, а осмысленный комментарий к коду. Поэтому ожидать больших значений метрики bleu не стоит.

После обучения модели её можно проверить на отдельной выборке:

Как можно увидеть, bleu-4 = 11, и это неплохо для такой задачи, даже с учётом того, что bleu в нашем случае распределена от 0 до 100.

Далее считаю получившиеся результаты:

Инициализирую функцию считывания из txt-файла:

И для удобства считаю всё в DataFrame:

Вывод 10 примеров кода, оригинальных комментариев и комментариев, сгенерированных моделью.

Теперь попробую субъективно сравнить оригинальный комментарий со сгенерированным по шкале от 1 до 5. Code — исходный код, true — исходный комментарий, pred — сгенерированный.

Пример 1:

Оценка: 1 — абсолютно непонятно, о чём идёт речь.

Пример 2:

Оценка: 4 — исходный комментарий абсолютно никак не отражает функциональность, в отличие от сгенерированного.

Пример 3:

Оценка: 5 — попадание в точку.

Пример 4:

Оценка: 3 — в целом, комментарии схожи, но вместо vertex используется shape, и в сгенерированном комментарии не отражено условие, которое прописано в оригинальном.

Пример 5:

Оценка: 2 — в исходном комментарии сказано, что генерируется новое буферное изображение определённого типа, в сгенерированном такие уточнения отсутствуют.

Пример 6:

Оценка: 1 — очень краткое и в то же время неверное описание.

Пример 7:

Оценка: 3 — в целом, суть похожа.

Пример 8:

Оценка: 3 — в целом, суть похожа.

Пример 9:

Оценка: 2 — не очень похоже на правду, но сгенерированный комментарий вполне осмысленный.

Пример 10:

Оценка: 1 — модель не смогла уловить суть.

Средняя субъективная оценка: (1+4+5+3+2+1+3+3+2+1)/10 = 2,5 — вполне неплохой результат для модели, которая училась на 1 % от общего объёма обучающих данных. В целом, суть сгенерированных комментариев понятна, но если у вас есть ресурсы, чтобы обучить модель более чем на 1 % набора, то вы можете улучшить результат.

Я показал, что после обучения модели даже на 1 % данных она выполняет свою цель и может вполне адекватно генерировать комментарии к коду. Также продемонстрирована предварительная обработка текста для языка Java. Если модель будет использоваться в исследовании целой кодовой базы, то лучше её всё же обучить на всех данных.

Также следует сказать, что если обучить модель на большем объёме, то её можно встроить в IDE (VisualStudio, PyCharm и т.д.) Подробнее об этом можно посмотреть здесь.