Fill-in-the-Middle: как модель дописывает код в середине файла

Клоуз-тест (cloze-test) - это упражнение, в котором некоторая часть текста скрывается (маскируется), и участнику необходимо заполнить эти пробелы учитывая, лексику и контекст. Слово cloze произошло от closure из гештальдпсихологии, а сам тест предложил Вилсон Тэйлор в далеком 1953 году [1]. В 2018 году похожий тест применили для обучения языковой модели BERT [2]. Если при классическом обучении LLM предсказывает следующий токен некоторой входной последовательности

token1, token2, token3, ... -> predict_next_token

то в cloze-подходе авторы статьи заставили BERT-модель предсказывать замаскированные токены, получая фрагменты контекста слева и справа, которые обычно называют PREFIX и SUFFIX

Подобный метод обучения моделей называют Fill-in-the-Middle (FIM). Эту идею можно применить и к тексту, содержащему код. Обычно при FIM-подходе во время обучения добавляют специальные токены, которые обозначают PREFIX, SUFFIX, MIDDLE (замаскированный код). На картинках эти токены представлены как |PREFIX|, |SUFFIX| и |MIDDLE| соответственно

В этом примере |END_OF_TEXT| также является специальным токеном, который обозначает конец генерации. Его можно воспринимать как один из стоп-сигналов: если модель сгенерировала его, клиент прекращает генерацию и возвращает ответ (без этого токена). На практике часто добавляют и другие стоп-условия.

В итоге, так как LLM просто продолжает входящую последовательность, отправляя на вход |PREFIX|{prefix}|SUFFIX|{suffix}|MIDDLE|, мы ожидаем получить {masked_code}|END_OF_TEXT|

Рассмотрим небольшой пример: пусть в некотором файле card.py мы замаскировали код Card("Q", "Hearts") и хотим, чтобы LLM его предсказала. Во время обучения модель учится предсказывать продолжение PREFIX так, чтобы сгенерированный код логично соединял PREFIX и SUFFIX. Она заметит в префиксе созданный датакласс Card, а из суффикса видно, что ниже ожидается объект card с полем rank, т.к. идет обращение к его полю card.rank. Составим промпт для LLM по принципу, описанному выше

В нашем примере мы можем получить примерно такой ответ

Обычно модели содержат и другие специальные токены, которые позволяют предоставлять модели дополнительный контекст: содержание других файлов, содержание pull request, репорты из GitHub/GitLab issues, и т.д.

Единого соглашения об именовании таких специальных токенов нет, поэтому мы используем обобщенный вид |{TOKEN_NAME}|, не привязываясь к конкретной модели. Вот небольшая таблица с примерами таких токенов для популярных кодинговых моделей

Токен |FILE_SEP|, как легко догадаться, служит для разграничения кода из разных файлов. Как видно из таблицы - не все модели во время обучения явно имеют специальный токен для этого. Например, модель Codestral в режиме multifile собирает несколько секций вида

+++++ path/to/file.ts  
{content}  
  
+++++ path/to/current.ts  
{prefix}

и после объединяет их в финальный шаблон вида

[SUFFIX]{suffix}[PREFIX]{all_files_plus_current_prefix}

Другими словами, все кодовые сниппеты добавляются к префиксу основного файла. Подробнее о том, как FIM используется при обучении кодовых моделей, можно прочитать в работах [3], [4], [5], [6], [7].

Тут важно отметить, что предсказывать замаскированный код можно не только с помощью FIM-обученной модели, но и с помощью обычной instruct LLM по типу GPT. Но обычно так не делают. FIM-модель обучена ровно на такую задачу, когда instruct-модель использует другой интерфейс: ей дают человеческую инструкцию, а она отвечает как ассистент. Примером может служить следующий промпт

Ожидается, что instruct-модель допишет только содержимое completion. Аналог |END_OF_TEXT| в этом запросе играют теги <COMPLETION>, </COMPLETION>

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