Autouzupełnianie w zintegrowanych środowiskach programistycznych (IDE) od dawna wspierane jest przez potęgę i precyzję odpowiednich funkcji. Jednak na arenę wkroczył nowy pretendent: Sztuczna Inteligencja (AI). Integracja AI w IDE staje się coraz bardziej powszechna, co skłania do ponownej oceny tradycyjnych narzędzi. Chociaż te starsze metody nadal działają, mogą już nie zapewniać najnowocześniejszej pomocy, jakiej oczekują programiści.
Ten artykuł, zaadaptowany z wpisu na blogu JetBrains, zagłębia się w szkolenie i możliwości Mellum, modelu stanowiącego podstawę chmurowego autouzupełniania. JetBrains, znany ze swoich IDE i edytorów kodu dostosowanych dla programistów, opracował Mellum, aby zapewnić szybkie i wydajne autouzupełnianie kodu. Mellum został zaprojektowany do efektywnej pracy na urządzeniu lokalnym, co wyróżnia go jako jeden z najszybszych i najbardziej kompaktowych modeli w swojej klasie. Koncepcyjnie podobnym projektem jest Microsoft Phi.
Wyzwania stojące przed Mellum
Inżynierowie JetBrains napotkali kilka krytycznych wyzwań w dążeniu do autouzupełniania opartego na sztucznej inteligencji:
- Szybkość i koszt: Konwencjonalne modele czatu okazały się niepraktyczne ze względu na wysokie koszty obliczeniowe i powolny czas reakcji. Modele te nie były również świadome technik specyficznych dla kodu, takich jak fill-in-the-middle (FIM) lub token healing.
- Formatowanie wyjścia: Flagowe modele czatu często generowały dane w niespójnych formatach, co utrudniało analizowanie odpowiedzi i bezproblemowe integrowanie ich z edytorem.
- Pochodzenie danych: Ustalenie pochodzenia danych szkoleniowych i ograniczenie potencjalnych problemów z naruszeniem praw autorskich stanowiło znaczną przeszkodę.
Mellum: Omówienie
Zespół programistów w JetBrains zdał sobie sprawę, że stworzenie własnego modelu jest optymalnym podejściem. Ich celem było zaprojektowanie modelu, który równoważyłby jakość, koszty wnioskowania i opóźnienia, a jednocześnie byłby szkolony na danych o jasnym pochodzeniu. Wstępne badania sugerowały, że model z około 4 miliardami parametrów może zapewnić kompetentne możliwości autouzupełniania dla szerokiego zakresu scenariuszy i użytkowników. Ponadto, trenując model wyłącznie na kodzie, mogli stworzyć specjalistyczny słownik tokenów, wolny od nieistotnych danych.
Proces szkolenia modelu Mellum składa się z trzech odrębnych etapów, z których każdy wnosi nową wiedzę i poprawia jakość generowanego kodu. Początkowa faza obejmuje podstawowe wstępne szkolenie na dużym korpusie pojedynczych plików. Drugi etap polega na udoskonaleniu modelu za pomocą mniejszego zestawu wyspecjalizowanych przykładów. Uczenie się przez wzmacnianie z informacją zwrotną od AI (RLAIF) jest wykorzystywane w trzecim etapie, aby dostosować model do specyficznych cech IDE i wyeliminować niepożądane dane wyjściowe.
Wstępne szkolenie
Aby uniknąć potencjalnych problemów związanych z niejednoznacznym pochodzeniem danych, model był trenowany od podstaw, co wymagało kompleksowego wprowadzenia do licznych języków, składni programowania, wzorców i podstawowych koncepcji.
Zbiór danych
Podstawowym źródłem danych do wstępnego szkolenia był TheStack. Ten zbiór danych zapewnia, że dane są zarówno zgodne z prawem, jak i praktycznie korzystne.
Proces wstępnego szkolenia
Podczas wstępnego szkolenia zbiór danych był próbkowany wielokrotnie, aby osiągnąć około 3 biliony tokenów. Użyto okna kontekstowego o długości 8192 tokenów, a zbiór danych podzielono na fragmenty o tym rozmiarze. Transformacja fill-in-the-middle (FIM) została zastosowana do połowy plików w każdym fragmencie, zachęcając model do rozważenia zarówno poprzedzającego, jak i następującego kodu. Technika ta ściśle naśladuje rzeczywiste scenariusze generowania kodu.
Faza wstępnego szkolenia została przeprowadzona na klastrze szesnastu węzłów, z których każdy był wyposażony w osiem procesorów graficznych H100. Ten etap zajął około 15 dni, w wyniku czego powstał 4-miliardowy model Mellum-base.
Wstępne szkolenie tworzy ogólny model autouzupełniania kodu z rozległą wiedzą na temat wielu języków programowania. Jednak na tym etapie model jest szkolony tylko do przewidywania następnego tokena w losowo wybranym segmencie pliku. Bez dodatkowego kontekstu model nie ma świadomości struktury kodu i nie ma mechanizmu do określania, kiedy przestać generować kod.
Etap dostrajania ma na celu wyeliminowanie tych ograniczeń.
Dostrajanie uwzględniające kontekst
Ulepszone Fill-in-the-Middle
W przeciwieństwie do wstępnego szkolenia, gdzie fragmenty kodu są wybierane losowo do przewidywania, dostrajanie koncentruje się na segmentowaniu kodu w bardziej znaczący sposób, ucząc model wyodrębniania fragmentów kodu, które występują “na wolności”.
Specjalistyczne przykłady
W praktyce autouzupełnianie kodu wymaga zrozumienia otaczających plików i szerszych kontekstów, potencjalnie obejmujących całe projekty.
Do przetwarzania wstępnego danych firma uruchomiła wewnętrzny projekt o nazwie kodowej Code Engine: wieloplatformowy zestaw SDK i zestaw narzędzi konsolowych opracowanych w celu budowania kontekstu bezpośrednio ze zwykłych plików bez konieczności pełnego indeksowania projektu. Ten zestaw SDK został wdrożony w wewnętrznym klastrze MapReduce i wykorzystany do przetwarzania tysięcy publicznych repozytoriów, generując wiele przydatnych przykładów do szkolenia w rozsądnym czasie.
Znalezienie właściwych algorytmów wymagało pewnych prób i błędów.
Dostrajanie pod kątem określonych języków
Małe modele mogą w dużym stopniu korzystać ze specjalizacji dla określonych języków. Chociaż model podstawowy jest trenowany w ponad 80 językach, większość użytkowników zazwyczaj pracuje tylko z jednym lub dwoma. Aby rozwiązać ten problem, JetBrains stworzył kilka wyspecjalizowanych modeli:
- mellum-all: Obsługuje większość języków i dialektów dostępnych w środowiskach IDE JetBrains, ale jakość autouzupełniania jest niższa niż w modelach specjalistycznych.
- mellum-python: Specjalizuje się w Pythonie i Jupyter.
- mellum-kotlin: Specjalizuje się w Javie i Kotlinie.
- mellum-web: Specjalizuje się w technologiach internetowych.
Ostatni krok: RLAIF
Wreszcie należy rozwiązać przypadki, w których cele uczenia się nie odpowiadają oczekiwaniom użytkowników. Dodatkowa faza szkolenia, RLAIF — Reinforcement Learning with AI Feedback, służy do rozwiązywania takich problemów.
Model uczy się na podstawie interakcji z użytkownikiem i rozumie, jak lepiej odzwierciedlać preferencje użytkownika.
Takie podejście nie tylko poprawia ogólny wynik jakości, ale także zmniejsza liczbę irytujących artefaktów generowania.
Jak dobry jest Mellum?
Model działa wyjątkowo dobrze, biorąc pod uwagę jego rozmiar. Oto jak został oceniony:
- Po pierwsze, model został oceniony na wewnętrznym benchmarku o nazwie kodowej “JetBrains BigCode”.
- Następnie został przetestowany na znanych publicznych benchmarkach, takich jak SAFIM.
- Wreszcie zebrano statystyki użytkowania funkcji i obliczono metryki użytkowników.
Ocena offline
Gromadzenie danych jest zadaniem złożonym, ale stworzenie dobrej metryki, która porównuje oryginalną sugestię z nową propozycją sieci neuronowej, jest jeszcze trudniejsze. Przeprowadziliśmy małe badanie i ostatecznie zdecydowaliśmy się na połączenie dwóch podstawowych metryk:
EM:
- Exact Match to bardzo popularny pomysł.
- Prognoza jest uważana za dobrą, jeśli pierwsza linia uzupełnienia pasuje do pierwszej linii oryginału, z minimalnym przetwarzaniem wstępnym.
KK:
- Metryka została nazwana na cześć jej autorów.
- Liczba proponowanych linii z oryginału podzielona przez liczbę linii w proponowanym uzupełnieniu.
JetBrains BigCode
Model został oceniony na zestawie danych benchmarkowych uzyskanym za pomocą wewnętrznego narzędzia JetBrains BigCode.
Utrzymując pełną kontrolę nad naszym zbiorem danych, zamiast polegać na publicznych benchmarkach, możemy niezawodnie oceniać jakość modelu dla różnych stylów i praktyk kodowania.
Wyniki naszej oceny JetBrains BigCode pokazują jakość porównywalną z popularnymi modelami, ale Mellum jest mniejszy i wydajniejszy.
Jakość jednowierszowych sugestii (metryka EM)
Publiczne benchmarki
Model został oceniony nie tylko na wewnętrznym zestawie danych, ale także na różnych publicznych benchmarkach, takich jak wielojęzyczny benchmark SAFIM (syntax-aware fill in the middle).
Ocena online
Główna metryka nazywa się współczynnikiem uzupełnionego kodu (RoCC). Jest on zdefiniowany jako stosunek znaków kodu napisanych przy użyciu autouzupełniania kodu do całkowitej ilości kodu w edytorze.
Inną ważną metryką jest współczynnik akceptacji (AR), który jest obliczany jako liczba zaakceptowanych sugestii podzielona przez liczbę wszystkich wyświetlonych sugestii.
To była złożona podróż, ale specjaliści JetBrains ukończyli ją z godnością. W rezultacie uzyskano jeden model ogólny i kilka wyspecjalizowanych, które są dostępne za pośrednictwem platformy JetBrains AI. Obecnie z powodzeniem działają w JetBrains AI Assistant.
Co dalej?
- Inżynierowie JetBrains pracują obecnie nad modelem dla języków programowania internetowego. Może on stać się publicznie dostępny w niedalekiej przyszłości.
- Planowane jest jednoczesne zwiększenie zarówno liczby parametrów, jak i różnorodności danych. W kodowaniu jest wiele różnych zadań — Mellum będzie mógł je również wykonywać. Wydajność usługi jest nadal kluczową metryką, więc rozbudowa modelu będzie w rozsądnych granicach.