De wereld van Integrated Development Environments (IDEs) wordt al lange tijd ondersteund door de kracht en precisie van autocompletion-functies. Er is echter een nieuwe speler op het toneel verschenen: Artificial Intelligence (AI). De integratie van AI in IDEs wordt steeds gebruikelijker, wat leidt tot een herziening van traditionele tools. Hoewel deze oudere methoden nog steeds functioneren, bieden ze mogelijk niet langer de geavanceerde hulp die ontwikkelaars nu verwachten.
Dit artikel is een bewerking van een blogpost van JetBrains en gaat dieper in op de training en mogelijkheden van Mellum, een model dat ten grondslag ligt aan cloud-gebaseerde autocompletion. JetBrains, bekend om zijn IDEs en code-editors die op maat zijn gemaakt voor programmeurs, heeft Mellum ontwikkeld om snelle en efficiënte code-autocompletion te bieden. Mellum is ontworpen om effectief te werken op een lokaal apparaat, waardoor het zich onderscheidt als een van de snelste en meest compacte modellen in zijn klasse. Een conceptueel vergelijkbaar project is Microsoft Phi.
Het Aanpakken van de Voor de Hand Liggende Uitdagingen
JetBrains-engineers stuitten op een aantal cruciale uitdagingen in hun streven naar AI-gestuurde autocompletion:
- Snelheid en Kosten: Conventionele chatmodellen bleken onpraktisch vanwege hun hoge rekenkosten en trage reactietijden. Deze modellen hadden ook geen kennis van codespecifieke technieken zoals fill-in-the-middle (FIM) of token healing.
- Output Formatting: Flagship chatmodellen genereerden vaak gegevens in inconsistente formaten, waardoor het moeilijk was om de reacties te parseren en naadloos in de editor te integreren.
- Data Provenance: Het bepalen van de oorsprong van de trainingsgegevens en het beperken van mogelijke schendingen van het auteursrecht vormden een aanzienlijke hindernis.
Mellum: Een Overzicht
Het ontwikkelingsteam van JetBrains realiseerde zich dat het creëren van hun eigen model de optimale aanpak was. Hun doel was om een model te ontwerpen dat kwaliteit, inference-kosten en latency in evenwicht bracht, terwijl het werd getraind op gegevens met een duidelijke herkomst. Voorlopig onderzoek suggereerde dat een model met ongeveer 4 miljard parameters competente autocompletion-mogelijkheden zou kunnen bieden voor een breed scala aan scenario’s en gebruikers. Bovendien konden ze, door het model uitsluitend op code te trainen, een gespecialiseerde vocabulaire van tokens vaststellen, vrij van irrelevante gegevens.
Het trainingsproces voor het Mellum-model bestaat uit drie verschillende fasen, die elk nieuwe kennis bijdragen en de kwaliteit van de gegenereerde code verbeteren. De initiële fase omvat basis pre-training op een grote corpus van individuele bestanden. De tweede fase bestaat uit het verfijnen van het model met een kleinere set gespecialiseerde voorbeelden. Reinforcement Learning with AI Feedback (RLAIF) wordt in de derde fase gebruikt om het model aan te passen aan IDE-specifieke kenmerken en ongewenste output te elimineren.
Pre-Training
Om potentiële problemen met betrekking tot dubbelzinnige gegevensherkomst te voorkomen, werd het model vanaf de grond af getraind, waarbij een uitgebreide introductie tot talloze talen, programmeersyntaxis, patronen en kernconcepten nodig was.
Dataset
De primaire gegevensbron voor pre-training was TheStack. Deze dataset zorgt ervoor dat de gegevens zowel juridisch correct als praktisch nuttig zijn.
Pre-Training Proces
Tijdens de pre-training werd de dataset meerdere keren gesampled om ongeveer 3 biljoen tokens te bereiken. Er werd een context window van 8192 tokens gebruikt, waarbij de dataset werd opgesplitst in fragmenten van deze grootte. De fill-in-the-middle (FIM) transformatie werd toegepast op de helft van de bestanden in elk fragment, waardoor het model werd aangemoedigd om zowel de voorgaande als de volgende code te overwegen. Deze techniek bootst real-world code generatie scenario’s nauwkeurig na.
De pre-training fase werd uitgevoerd op een cluster van zestien nodes, elk uitgerust met acht H100 GPUs. Deze fase duurde ongeveer 15 dagen om te voltooien, wat resulteerde in het 4-miljard-parameter Mellum-base model.
Pre-training creëert een general-purpose code autocompletion model met uitgebreide kennis van vele programmeertalen. In dit stadium is het model echter alleen getraind om het volgende token in een willekeurig geselecteerd bestandsegment te voorspellen. Zonder extra context heeft het model geen kennis van de codestructuur en heeft het geen mechanisme om te bepalen wanneer het moet stoppen met het genereren van code.
De fine-tuning fase is ontworpen om deze beperkingen aan te pakken.
Context-Aware Fine-Tuning
Verbeterde Fill-in-the-Middle
In tegenstelling tot pre-training, waarbij codefragmenten willekeurig worden geselecteerd voor voorspelling, concentreert fine-tuning zich op het segmenteren van code op een meer betekenisvolle manier, waardoor het model leert om codefragmenten te extraheren die “in het wild” voorkomen.
Gespecialiseerde Voorbeelden
In de praktijk vereist code autocompletion het begrijpen van omliggende bestanden en bredere contexten, mogelijk inclusief hele projecten.
Voor data preprocessing lanceerde het bedrijf een intern project met de codenaam Code Engine: een cross-platform SDK en een set console utilities ontwikkeld om context rechtstreeks uit gewone bestanden te bouwen zonder volledige projectindexering vereist. Deze SDK werd ingezet op een intern MapReduce-cluster en gebruikt om duizenden openbare repositories te verwerken, waardoor veel nuttige voorbeelden werden gegenereerd voor training binnen een redelijk tijdsbestek.
Het vinden van de juiste algoritmen vereiste enig vallen en opstaan.
Tuning voor Specifieke Talen
Kleine modellen kunnen enorm profiteren van specialisatie voor specifieke talen. Hoewel het basismodel is getraind op meer dan 80 talen, werken de meeste gebruikers doorgaans met slechts één of twee. Om dit aan te pakken, creëerde JetBrains meerdere gespecialiseerde modellen:
- mellum-all: Ondersteunt de meeste talen en dialecten die beschikbaar zijn in JetBrains IDEs, maar de autocompletion-kwaliteit is lager dan bij gespecialiseerde modellen.
- mellum-python: Gespecialiseerd in Python en Jupyter.
- mellum-kotlin: Gespecialiseerd in Java en Kotlin.
- mellum-web: Gespecialiseerd in webtechnologieën.
De Laatste Stap: RLAIF
Ten slotte moeten gevallen waarin de leerdoelen niet overeenkomen met de verwachtingen van de gebruiker worden opgelost. Een extra trainingsfase, RLAIF — Reinforcement Learning with AI Feedback, wordt gebruikt om dergelijke problemen op te lossen.
Het model leert van gebruikersinteracties en begrijpt hoe het de voorkeuren van de gebruiker beter kan weerspiegelen.
Deze aanpak verbetert niet alleen de algehele kwaliteitsscore, maar vermindert ook het aantal irritante generatie-artefacten.
Hoe Goed Is Mellum?
Het model presteert uitzonderlijk goed voor zijn omvang. Hier is hoe het werd geëvalueerd:
- Eerst werd het model geëvalueerd op een interne benchmark met de codenaam “JetBrains BigCode”.
- Vervolgens werd het getest op bekende openbare benchmarks zoals SAFIM.
- Ten slotte werden gebruiksstatistieken voor functies verzameld en werden gebruikersstatistieken berekend.
Offline Evaluatie
Het verzamelen van gegevens is een complexe taak, maar het creëren van een goede metric die de originele suggestie vergelijkt met de nieuwe die door het neurale netwerk wordt voorgesteld, is nog uitdagender. We voerden een kleine studie uit en kwamen uiteindelijk uit op een combinatie van twee primaire metrics:
EM:
- Exact Match is een erg populair idee.
- Een voorspelling wordt als goed beschouwd als de eerste regel van de completion overeenkomt met de eerste regel van het origineel, met minimale preprocessing.
KK:
- De metric is vernoemd naar de auteurs.
- Het aantal voorgestelde regels uit het origineel gedeeld door het aantal regels in de voorgestelde completion.
JetBrains BigCode
Het model werd geëvalueerd aan de hand van een benchmark dataset verkregen met behulp van de interne JetBrains BigCode tool.
Door de volledige controle over onze dataset te behouden in plaats van te vertrouwen op openbare benchmarks, wordt het mogelijk om de modelkwaliteit betrouwbaar te evalueren voor verschillende codeerstijlen en -praktijken.
De resultaten van onze JetBrains BigCode-evaluatie laten een kwaliteit zien die vergelijkbaar is met populaire modellen, maar Mellum is kleiner en efficiënter.
Kwaliteit van single-line suggesties (EM metric)
Openbare Benchmarks
Het model werd niet alleen geëvalueerd op de interne dataset, maar ook op verschillende openbare benchmarks, zoals de meertalige benchmark SAFIM (syntax-aware fill in the middle).
Online Evaluatie
De belangrijkste metric wordt de ratio of completed code (RoCC) genoemd. Het wordt gedefinieerd als de verhouding van code characters geschreven met behulp van code autocompletion tot de totale hoeveelheid code in de editor.
Een andere belangrijke metric is de acceptance rate (AR), die wordt berekend als het aantal geaccepteerde suggesties gedeeld door het aantal van alle getoonde suggesties.
Dit was een complexe reis, maar JetBrains-specialisten hebben deze met waardigheid voltooid. Uiteindelijk werden een algemeen en verschillende gespecialiseerde modellen verkregen, die beschikbaar zijn via het JetBrains AI-platform. Ze werken nu succesvol in JetBrains AI Assistant.
Wat is de volgende stap?
- JetBrains-engineers werken momenteel aan een model voor web development talen. Het kan in de nabije toekomst openbaar beschikbaar worden.
- Er zijn plannen om tegelijkertijd zowel het aantal parameters als de diversiteit van de gegevens te vergroten. Er zijn veel verschillende taken in coding — Mellum zal ze ook kunnen uitvoeren. Service performance is nog steeds een belangrijke metric, dus de uitbreiding van het model zal binnen redelijke grenzen blijven.