Mellum: Un Modello Veloce e Miniaturizzato per il Completamento del Codice nel Tuo Editor
Il regno degli Integrated Development Environments (IDE) è stato a lungo sostenuto dalla potenza e dalla precisione delle funzionalità di autocompletamento. Tuttavia, un nuovo contendente è entrato nell’arena: l’Intelligenza Artificiale (AI). L’integrazione dell’AI negli IDE sta diventando sempre più ubiquitaria, spingendo a una rivalutazione degli strumenti tradizionali. Mentre questi metodi più vecchi funzionano ancora, potrebbero non fornire più l’assistenza all’avanguardia che gli sviluppatori si aspettano ora.
Questo articolo è adattato da un post del blog di JetBrains e approfondisce la formazione e le capacità di Mellum, un modello alla base dell’autocompletamento basato su cloud. JetBrains, rinomata per i suoi IDE e editor di codice su misura per i programmatori, ha sviluppato Mellum per fornire un autocompletamento del codice rapido ed efficiente. Mellum è progettato per operare efficacemente su un dispositivo locale, distinguendosi come uno dei modelli più veloci e compatti della sua classe. Un progetto concettualmente simile è Microsoft Phi.
Affrontare le Ovvie Sfide
Gli ingegneri di JetBrains hanno incontrato diverse sfide critiche nella loro ricerca dell’autocompletamento basato sull’AI:
- Velocità e Costo: I modelli di chat convenzionali si sono dimostrati impraticabili a causa dei loro elevati costi computazionali e dei tempi di risposta lenti. Questi modelli mancavano anche di consapevolezza delle tecniche specifiche del codice come fill-in-the-middle (FIM) o token healing.
- Formattazione dell’Output: I modelli di chat di punta spesso generavano dati in formati incoerenti, rendendo difficile l’analisi delle risposte e la loro integrazione perfetta nell’editor.
- Provenienza dei Dati: Determinare l’origine dei dati di addestramento e mitigare potenziali problemi di violazione del copyright ha rappresentato un ostacolo significativo.
Mellum: Una Panoramica
Il team di sviluppo di JetBrains si è reso conto che creare il proprio modello era l’approccio ottimale. Il loro obiettivo era progettare un modello che bilanciasse qualità, costi di inferenza e latenza, pur essendo addestrato su dati con una chiara provenienza. La ricerca preliminare ha suggerito che un modello con circa 4 miliardi di parametri potrebbe fornire capacità di autocompletamento competenti per una vasta gamma di scenari e utenti. Inoltre, addestrando esclusivamente il modello sul codice, potrebbero stabilire un vocabolario specializzato di token, privo di dati irrilevanti.
Il processo di addestramento per il modello Mellum consiste in tre fasi distinte, ognuna delle quali contribuisce con nuove conoscenze e migliora la qualità del codice generato. La fase iniziale prevede il pre-addestramento di base su un ampio corpus di singoli file. La seconda fase consiste nel perfezionare il modello con un set più piccolo di esempi specializzati. Il Reinforcement Learning with AI Feedback (RLAIF) viene impiegato nella terza fase per adattare il modello alle caratteristiche specifiche dell’IDE ed eliminare l’output indesiderato.
Pre-Addestramento
Per evitare potenziali problemi relativi alle origini ambigue dei dati, il modello è stato addestrato da zero, richiedendo un’introduzione completa a numerose lingue, sintassi di programmazione, modelli e concetti di base.
Dataset
La principale fonte di dati per il pre-addestramento è stata TheStack. Questo dataset garantisce che i dati siano sia legalmente validi che praticamente vantaggiosi.
Processo di Pre-Addestramento
Durante il pre-addestramento, il dataset è stato campionato più volte per raggiungere circa 3 trilioni di token. È stata utilizzata una finestra di contesto di 8192 token, con il dataset suddiviso in frammenti di questa dimensione. La trasformazione fill-in-the-middle (FIM) è stata applicata alla metà dei file in ogni frammento, incoraggiando il modello a considerare sia il codice precedente che quello successivo. Questa tecnica imita da vicino gli scenari di generazione di codice nel mondo reale.
La fase di pre-addestramento è stata condotta su un cluster di sedici nodi, ciascuno dotato di otto GPU H100. Questa fase ha richiesto circa 15 giorni per essere completata, risultando nel modello Mellum-base con 4 miliardi di parametri.
Il pre-addestramento crea un modello di autocompletamento del codice generico con un’ampia conoscenza di molti linguaggi di programmazione. Tuttavia, in questa fase, il modello è addestrato solo a prevedere il token successivo in un segmento di file selezionato casualmente. Senza un contesto aggiuntivo, il modello manca di consapevolezza della struttura del codice e non ha alcun meccanismo per determinare quando smettere di generare codice.
La fase di fine-tuning è progettata per affrontare queste limitazioni.
Fine-Tuning Consapevole del Contesto
Fill-in-the-Middle Migliorato
A differenza del pre-addestramento, dove i frammenti di codice vengono selezionati casualmente per la previsione, il fine-tuning si concentra sulla segmentazione del codice in un modo più significativo, insegnando al modello a estrarre frammenti di codice che si verificano ‘in natura’.
Esempi Specializzati
In pratica, l’autocompletamento del codice necessita della comprensione dei file circostanti e di contesti più ampi, possibilmente comprendendo interi progetti.
Per la pre-elaborazione dei dati, l’azienda ha lanciato un progetto interno con il nome in codice Code Engine: un SDK multipiattaforma e un set di utility della console sviluppati per costruire il contesto direttamente da file ordinari senza richiedere l’indicizzazione completa del progetto. Questo SDK è stato distribuito su un cluster MapReduce interno e utilizzato per elaborare migliaia di repository pubblici, generando molti esempi utili per l’addestramento in un lasso di tempo ragionevole.
Trovare gli algoritmi corretti ha richiesto alcuni tentativi ed errori.
Ottimizzazione per Linguaggi Specifici
I piccoli modelli possono trarre grande vantaggio dalla specializzazione per linguaggi specifici. Mentre il modello base è addestrato su oltre 80 lingue, la maggior parte degli utenti in genere lavora solo con uno o due. Per risolvere questo problema, JetBrains ha creato più modelli specializzati:
- mellum-all: Supporta la maggior parte delle lingue e dei dialetti disponibili negli IDE di JetBrains, ma la qualità dell’autocompletamento è inferiore rispetto ai modelli specializzati.
- mellum-python: È specializzato in Python e Jupyter.
- mellum-kotlin: È specializzato in Java e Kotlin.
- mellum-web: È specializzato in tecnologie web.
Il Passo Finale: RLAIF
Infine, i casi in cui gli obiettivi di apprendimento non corrispondono alle aspettative degli utenti devono essere risolti. Un’ulteriore fase di addestramento, RLAIF — Reinforcement Learning with AI Feedback, viene utilizzata per risolvere tali problemi.
Il modello apprende dalle interazioni degli utenti e comprende come riflettere meglio le preferenze degli utenti.
Questo approccio non solo migliora il punteggio di qualità complessivo, ma riduce anche il numero di artefatti di generazione fastidiosi.
Quanto è Buono Mellum?
Il modello si comporta eccezionalmente bene per le sue dimensioni. Ecco come è stato valutato:
- Innanzitutto, il modello è stato valutato su un benchmark interno con il nome in codice ‘JetBrains BigCode’.
- Quindi è stato testato su benchmark pubblici ben noti come SAFIM.
- Infine, sono state raccolte le statistiche di utilizzo per le funzionalità e sono state calcolate le metriche utente.
Valutazione Offline
La raccolta di dati è un compito complesso, ma la creazione di una buona metrica che confronti il suggerimento originale con quello nuovo proposto dalla rete neurale è ancora più impegnativa. Abbiamo condotto un piccolo studio e alla fine ci siamo stabiliti su una combinazione di due metriche principali:
EM:
- Exact Match è un’idea molto popolare.
- Una previsione è considerata buona se la prima riga del completamento corrisponde alla prima riga dell’originale, con una pre-elaborazione minima.
KK:
- La metrica è stata chiamata in onore dei suoi autori.
- Il numero di righe proposte dall’originale diviso per il numero di righe nel completamento proposto.
JetBrains BigCode
Il modello è stato valutato rispetto a un dataset di benchmark ottenuto utilizzando lo strumento interno JetBrains BigCode.
Mantenendo il pieno controllo sul nostro dataset piuttosto che affidarci a benchmark pubblici, diventa possibile valutare in modo affidabile la qualità del modello per vari stili e pratiche di codifica.
I risultati della nostra valutazione JetBrains BigCode mostrano una qualità alla pari con i modelli popolari, ma Mellum è più piccolo ed efficiente.
Qualità dei suggerimenti a riga singola (metrica EM)
Benchmark Pubblici
Il modello è stato valutato non solo sul dataset interno, ma anche su vari benchmark pubblici, come il benchmark multilingue SAFIM (syntax-aware fill in the middle).
Valutazione Online
La metrica principale è chiamata rapporto di codice completato (RoCC). È definito come il rapporto tra i caratteri di codice scritti utilizzando l’autocompletamento del codice e la quantità totale di codice nell’editor.
Un’altra metrica importante è il tasso di accettazione (AR), che viene calcolato come il numero di suggerimenti accettati diviso per il numero di tutti i suggerimenti mostrati.
Questo è stato un viaggio complesso, ma gli specialisti di JetBrains lo hanno completato con dignità. Alla fine, sono stati ottenuti un modello generale e diversi modelli specializzati, che sono disponibili tramite la piattaforma JetBrains AI. Ora stanno lavorando con successo in JetBrains AI Assistant.
Cosa c’è di nuovo?
- Gli ingegneri di JetBrains stanno attualmente lavorando a un modello per i linguaggi di sviluppo web. Potrebbe diventare disponibile pubblicamente nel prossimo futuro.
- Ci sono piani per aumentare contemporaneamente sia il numero di parametri che la diversità dei dati. Ci sono molti compiti diversi nella codifica: Mellum sarà in grado di svolgerli anche lui. Le prestazioni del servizio sono ancora una metrica chiave, quindi l’espansione del modello sarà entro limiti ragionevoli.