Mellum: Schnelles Code-Vervollständigungsmodell

In der Welt der integrierten Entwicklungsumgebungen (IDEs) hat die Leistungsfähigkeit und Präzision von Autovervollständigungsfunktionen lange eine wichtige Rolle gespielt. Allerdings ist ein neuer Mitbewerber in das Geschehen eingetreten: die künstliche Intelligenz (KI). Die Integration von KI in IDEs wird immer allgegenwärtiger und führt zu einer Neubewertung traditioneller Werkzeuge. Obwohl diese älteren Methoden weiterhin funktionieren, bieten sie möglicherweise nicht mehr die hochmoderne Unterstützung, die Entwickler heute erwarten.

Dieser Artikel basiert auf einem Blog-Beitrag von JetBrains und befasst sich mit dem Training und den Fähigkeiten von Mellum, einem Modell, das die Cloud-basierte Autovervollständigung unterstützt. JetBrains, bekannt für seine IDEs und Code-Editoren, die auf Programmierer zugeschnitten sind, hat Mellum entwickelt, um eine schnelle und effiziente Code-Autovervollständigung zu ermöglichen. Mellum ist so konzipiert, dass es effektiv auf einem lokalen Gerät arbeitet, was es als eines der schnellsten und kompaktesten Modelle seiner Klasse auszeichnet. Ein konzeptionell ähnliches Projekt ist Microsoft Phi.

Bewältigung der offensichtlichen Herausforderungen

Die Ingenieure von JetBrains stießen bei ihrer Suche nach KI-gestützter Autovervollständigung auf mehrere kritische Herausforderungen:

  • Geschwindigkeit und Kosten: Herkömmliche Chatmodelle erwiesen sich aufgrund ihrer hohen Rechenkosten und langsamen Reaktionszeiten als unpraktisch. Diesen Modellen fehlte auch das Bewusstsein für codespezifische Techniken wie Fill-in-the-Middle (FIM) oder Token Healing.
  • Ausgabeformatierung: Flaggschiff-Chatmodelle generierten Daten oft in inkonsistenten Formaten, was es erschwerte, die Antworten zu analysieren und sie nahtlos in den Editor zu integrieren.
  • Datenherkunft: Die Bestimmung des Ursprungs der Trainingsdaten und die Minderung potenzieller Urheberrechtsverletzungen stellten eine erhebliche Hürde dar.

Mellum: Ein Überblick

Das Entwicklungsteam von JetBrains erkannte, dass die Erstellung eines eigenen Modells der optimale Ansatz war. Ihr Ziel war es, ein Modell zu entwerfen, das Qualität, Inferenzkosten und Latenz ausbalanciert und gleichzeitig auf Daten mit klarer Herkunft trainiert wird. Vorläufige Untersuchungen deuteten darauf hin, dass ein Modell mit etwa 4 Milliarden Parametern kompetente Autovervollständigungsfunktionen für eine Vielzahl von Szenarien und Benutzern bieten könnte. Darüber hinaus könnten sie durch ausschließliches Training des Modells mit Code ein spezielles Vokabular von Token erstellen, das frei von irrelevanten Daten ist.

Der Trainingsprozess für das Mellum-Modell besteht aus drei verschiedenen Phasen, die jeweils neues Wissen beitragen und die Qualität des generierten Codes verbessern. Die erste Phase umfasst ein grundlegendes Vortraining auf einem großen Korpus einzelner Dateien. Die zweite Phase besteht darin, das Modell mit einer kleineren Menge spezialisierter Beispiele zu verfeinern. Reinforcement Learning with AI Feedback (RLAIF) wird in der dritten Phase eingesetzt, um das Modell an IDE-spezifische Eigenschaften anzupassen und unerwünschte Ausgaben zu eliminieren.

Vortraining

Um potenzielle Probleme im Zusammenhang mit uneindeutigen Datenherkünften zu vermeiden, wurde das Modell von Grund auf neu trainiert, was eine umfassende Einführung in zahlreiche Sprachen, Programmiersyntax, Muster und Kernkonzepte erforderte.

Datensatz

Die primäre Datenquelle für das Vortraining war TheStack. Dieser Datensatz stellt sicher, dass die Daten sowohl rechtlich einwandfrei als auch praktisch vorteilhaft sind.

Vortrainingsprozess

Während des Vortrainings wurde der Datensatz mehrmals gesampelt, um ungefähr 3 Billionen Token zu erreichen. Es wurde ein Kontextfenster von 8192 Token verwendet, wobei der Datensatz in Fragmente dieser Größe aufgeteilt wurde. Die Fill-in-the-Middle (FIM)-Transformation wurde auf die Hälfte der Dateien in jedem Fragment angewendet, um das Modell zu ermutigen, sowohl den vorhergehenden als auch den nachfolgenden Code zu berücksichtigen. Diese Technik ahmt reale Code-Generierungsszenarien genau nach.

Die Vortrainingsphase wurde auf einem Cluster von sechzehn Knoten durchgeführt, von denen jeder mit acht H100-GPUs ausgestattet war. Diese Phase dauerte ungefähr 15 Tage, was zum 4-Milliarden-Parameter-Mellum-Basismodell führte.

Das Vortraining erstellt ein Allzweck-Code-Autovervollständigungsmodell mit umfassenden Kenntnissen vieler Programmiersprachen. In dieser Phase wird das Modell jedoch nur trainiert, um das nächste Token in einem zufällig ausgewählten Dateisegment vorherzusagen. Ohne zusätzlichen Kontext fehlt dem Modell das Bewusstsein für die Codestruktur und es gibt keinen Mechanismus, um zu bestimmen, wann die Codegenerierung beendet werden soll.

Die Feinabstimmungsphase ist darauf ausgelegt, diese Einschränkungen zu beheben.

Kontextbezogene Feinabstimmung

Verbessertes Fill-in-the-Middle

Im Gegensatz zum Vortraining, bei dem Codefragmente zufällig zur Vorhersage ausgewählt werden, konzentriert sich die Feinabstimmung auf die Segmentierung von Code auf sinnvollere Weise und lehrt das Modell, Codefragmente zu extrahieren, die ‘in der Wildnis’ vorkommen.

Spezialisierte Beispiele

In der Praxis erfordert die Code-Autovervollständigung das Verständnis umliegender Dateien und breiterer Kontexte, möglicherweise ganzer Projekte.

Für die Datenvorverarbeitung startete das Unternehmen ein internes Projekt mit dem Codenamen Code Engine: ein plattformübergreifendes SDK und eine Reihe von Konsolenprogrammen, die entwickelt wurden, um Kontext direkt aus gewöhnlichen Dateien zu erstellen, ohne eine vollständige Projektindizierung zu erfordern. Dieses SDK wurde auf einem internen MapReduce-Cluster bereitgestellt und verwendet, um Tausende von öffentlichen Repositories zu verarbeiten, wodurch in einem angemessenen Zeitrahmen viele nützliche Beispiele für das Training generiert wurden.

Das Finden der richtigen Algorithmen erforderte einiges an Ausprobieren.

Optimierung für bestimmte Sprachen

Kleine Modelle können stark von der Spezialisierung für bestimmte Sprachen profitieren. Während das Basismodell mit über 80 Sprachen trainiert wird, arbeiten die meisten Benutzer in der Regel nur mit ein oder zwei Sprachen. Um dies zu beheben, hat JetBrains mehrere spezialisierte Modelle erstellt:

  • mellum-all: Unterstützt die meisten Sprachen und Dialekte, die in JetBrains IDEs verfügbar sind, aber die Autovervollständigungsqualität ist geringer als bei spezialisierten Modellen.
  • mellum-python: Spezialisiert auf Python und Jupyter.
  • mellum-kotlin: Spezialisiert auf Java und Kotlin.
  • mellum-web: Spezialisiert auf Webtechnologien.

Der letzte Schritt: RLAIF

Schließlich müssen Fälle gelöst werden, in denen die Lernziele nicht mit den Erwartungen der Benutzer übereinstimmen. Eine zusätzliche Trainingsphase, RLAIF – Reinforcement Learning with AI Feedback, wird verwendet, um solche Probleme zu lösen.

Das Modell lernt aus Benutzerinteraktionen und versteht, wie es die Benutzerpräferenzen besser widerspiegeln kann.

Dieser Ansatz verbessert nicht nur die Gesamtqualitätsbewertung, sondern reduziert auch die Anzahl der störenden Generierungsartefakte.

Wie gut ist Mellum?

Das Modell schneidet für seine Größe außergewöhnlich gut ab. So wurde es bewertet:

  • Zuerst wurde das Modell auf einem internen Benchmark mit dem Codenamen ‘JetBrains BigCode’ bewertet.
  • Dann wurde es auf bekannten öffentlichen Benchmarks wie SAFIM getestet.
  • Schließlich wurden Nutzungsstatistiken für Funktionen gesammelt und Benutzermetriken berechnet.

Offline-Bewertung

Das Sammeln von Daten ist eine komplexe Aufgabe, aber die Erstellung einer guten Metrik, die den ursprünglichen Vorschlag mit dem neuen vergleicht, der vom neuronalen Netzwerk vorgeschlagen wird, ist noch schwieriger. Wir haben eine kleine Studie durchgeführt und uns letztendlich für eine Kombination aus zwei primären Metriken entschieden:

EM:

  • Exact Match ist eine sehr beliebte Idee.
  • Eine Vorhersage gilt als gut, wenn die erste Zeile der Vervollständigung mit der ersten Zeile des Originals übereinstimmt, mit minimaler Vorverarbeitung.

KK:

  • Die Metrik wurde nach ihren Autoren benannt.
  • Die Anzahl der vorgeschlagenen Zeilen aus dem Original geteilt durch die Anzahl der Zeilen in der vorgeschlagenen Vervollständigung.

JetBrains BigCode

Das Modell wurde anhand eines Benchmark-Datensatzes bewertet, der mit dem internen JetBrains BigCode-Tool erstellt wurde.

Durch die Aufrechterhaltung der vollständigen Kontrolle über unseren Datensatz, anstatt uns auf öffentliche Benchmarks zu verlassen, wird es möglich, die Modellqualität für verschiedene Codierungsstile und -praktiken zuverlässig zu bewerten.

Die Ergebnisse unserer JetBrains BigCode-Bewertung zeigen eine Qualität, die mit der beliebter Modelle vergleichbar ist, aber Mellum ist kleiner und effizienter.

Qualität von einzeiligen Vorschlägen (EM-Metrik)

Öffentliche Benchmarks

Das Modell wurde nicht nur auf dem internen Datensatz bewertet, sondern auch auf verschiedenen öffentlichen Benchmarks, wie z. B. dem mehrsprachigen Benchmark SAFIM (syntax-aware fill in the middle).

Online-Bewertung

Die Hauptmetrik wird als Ratio of Completed Code (RoCC) bezeichnet. Sie ist definiert als das Verhältnis der Codezeichen, die mit der Code-Autovervollständigung geschrieben wurden, zur Gesamtmenge des Codes im Editor.

Eine weitere wichtige Metrik ist die Acceptance Rate (AR), die als die Anzahl der akzeptierten Vorschläge geteilt durch die Anzahl aller angezeigten Vorschläge berechnet wird.

Dies war eine komplexe Reise, aber die Spezialisten von JetBrains haben sie mit Würde abgeschlossen. Am Ende wurden ein allgemeines und mehrere spezialisierte Modelle erhalten, die über die JetBrains AI-Plattform verfügbar sind. Sie arbeiten jetzt erfolgreich im JetBrains AI Assistant.

Was kommt als Nächstes?

  • Die Ingenieure von JetBrains arbeiten derzeit an einem Modell für Webentwicklungs sprachen. Es könnte in naher Zukunft öffentlich verfügbar werden.
  • Es gibt Pläne, gleichzeitig sowohl die Anzahl der Parameter als auch die Vielfalt der Daten zu erhöhen. Es gibt viele verschiedene Aufgaben beim Codieren – Mellum wird auch in der Lage sein, diese auszuführen. Die Serviceleistung ist nach wie vor eine Schlüsselmetrik, daher wird die Erweiterung des Modells in einem angemessenen Rahmen erfolgen.