Mellum: Бърз модел за довършване на код

В света на интегрираните среди за разработка (IDE) отдавна се разчита на мощта и прецизността на функциите за автоматично довършване. Въпреки това, нов претендент навлезе на сцената: Изкуственият интелект (AI). Интеграцията на AI в IDE става все по-разпространена, което налага преоценка на традиционните инструменти. Въпреки че тези по-стари методи все още работят, те може вече да не предоставят най-съвременната помощ, която разработчиците очакват.

Тази статия е адаптирана от публикация в блог на JetBrains и се задълбочава в обучението и възможностите на Mellum, модел, който е в основата на облачното автоматично довършване. JetBrains, известни със своите IDE и редактори на код, пригодени за програмисти, разработиха Mellum, за да осигурят бързо и ефикасно автоматично довършване на код. Mellum е проектиран да работи ефективно на локално устройство, което го отличава като един от най-бързите и компактни модели в своя клас. Концептуално подобен проект е Microsoft Phi.

Решаване на очевидните предизвикателства

Инженерите на JetBrains се сблъскаха с няколко критични предизвикателства в стремежа си към автоматично довършване, задвижвано от AI:

  • Скорост и цена: Конвенционалните чат модели се оказаха непрактични поради високите си изчислителни разходи и бавното време за реакция. Тези модели също така не познаваха специфични за кода техники, като fill-in-the-middle (FIM) или token healing.
  • Форматиране на изхода: Водещите чат модели често генерираха данни в непоследователни формати, което затрудняваше анализирането на отговорите и безпроблемното им интегриране в редактора.
  • Произход на данните: Определянето на произхода на данните за обучение и смекчаването на потенциални проблеми с нарушаване на авторските права представляваше значителна пречка.

Mellum: Общ преглед

Екипът за разработка в JetBrains осъзна, че създаването на собствен модел е оптималният подход. Тяхната цел беше да проектират модел, който балансира качество, разходи за извод и латентност, като същевременно е обучен на данни с ясен произход. Предварителните изследвания показват, че модел с около 4 милиарда параметри може да осигури компетентни възможности за автоматично довършване за широк спектър от сценарии и потребители. Освен това, като обучат модела изключително върху код, те могат да установят специализиран речник от токени, свободен от неподходящи данни.

Процесът на обучение за модела Mellum се състои от три отделни етапа, всеки от които допринася с нови знания и подобрява качеството на генерирания код. Първоначалната фаза включва основно предварително обучение върху голям корпус от отделни файлове. Вторият етап се състои в усъвършенстване на модела с по-малък набор от специализирани примери. Reinforcement Learning with AI Feedback (RLAIF) се използва в третия етап, за да се адаптира модела към специфичните за IDE характеристики и да се елиминира нежеланият изход.

Предварително обучение

За да се избегнат потенциални проблеми, свързани с неясен произход на данните, моделът беше обучен от нулата, което изискваше цялостно въведение в многобройни езици, програмен синтаксис, модели и основни концепции.

Набор от данни

Основният източник на данни за предварително обучение беше TheStack. Този набор от данни гарантира, че данните са едновременно правно обосновани и практически полезни.

Процес на предварително обучение

По време на предварителното обучение наборът от данни беше семплиран многократно, за да достигне приблизително 3 трилиона токена. Беше използван контекстов прозорец от 8192 токена, като наборът от данни беше разделен на фрагменти с този размер. Трансформацията fill-in-the-middle (FIM) беше приложена към половината от файловете във всеки фрагмент, насърчавайки модела да обмисли както предходния, така и последващия код. Тази техника тясно имитира реални сценарии за генериране на код.

Фазата на предварително обучение беше проведена на клъстер от шестнадесет възела, всеки оборудван с осем H100 GPU. Този етап отне приблизително 15 дни, което доведе до Mellum-base модела с 4 милиарда параметри.

Предварителното обучение създава модел за автоматично довършване на код с общо предназначение с обширни познания за много програмни езици. Въпреки това, на този етап, моделът е обучен само да предсказва следващия токен в произволно избран файлов сегмент. Без допълнителен контекст, моделът няма представа за структурата на кода и няма механизъм да определи кога да спре да генерира код.

Етапът на фината настройка е предназначен да адресира тези ограничения.

Контекстно-зависима фина настройка

Подобрен Fill-in-the-Middle

За разлика от предварителното обучение, където кодовите фрагменти се избират на случаен принцип за предсказване, фината настройка се концентрира върху сегментиране на кода по по-смислен начин, като учи модела да извлича кодови фрагменти, които се срещат ‘в дивата природа’.

Специализирани примери

На практика автоматичното довършване на код налага разбиране на заобикалящите файлове и по-широки контексти, евентуално обхващащи цели проекти.

За предварителна обработка на данни компанията стартира вътрешен проект, кодиран като Code Engine: кросплатформен SDK и набор от конзолни помощни програми, разработени за изграждане на контекст директно от обикновени файлове, без да се изисква пълно индексиране на проекта. Този SDK беше разгърнат на вътрешен MapReduce клъстер и използван за обработка на хиляди публични хранилища, генерирайки много полезни примери за обучение в разумен срок.

Намирането на правилните алгоритми изискваше известна степен на опити и грешки.

Настройка за специфични езици

Малките модели могат да се възползват значително от специализация за специфични езици. Докато базовият модел е обучен на над 80 езика, повечето потребители обикновено работят само с един или два. За да се справи с това, JetBrains създаде множество специализирани модели:

  • mellum-all: Поддържа повечето езици и диалекти, налични в JetBrains IDE, но качеството на автоматичното довършване е по-ниско от специализираните модели.
  • 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 ще може да ги изпълнява също. Производителността на услугата все още е ключов показател, така че разширяването на модела ще бъде в разумни граници.