L’intégration de l’intelligence artificielle (IA) dans les environnements de développement intégrés (IDE) est devenue de plus en plus courante, ce qui nécessite une réévaluation des outils traditionnels. Bien que ces anciennes méthodes fonctionnent toujours, elles ne fournissent peut-être plus l’assistance de pointe que les développeurs attendent désormais. Cet article, adapté d’un article de blog de JetBrains, se penche sur la formation et les capacités de Mellum, un modèle qui sous-tend la complétion automatique basée sur le cloud. JetBrains, réputé pour ses IDE et ses éditeurs de code adaptés aux programmeurs, a développé Mellum pour fournir une complétion de code rapide et efficace. Mellum est conçu pour fonctionner efficacement sur un appareil local, ce qui le distingue comme l’un des modèles les plus rapides et les plus compacts de sa catégorie. Un projet conceptuellement similaire est Microsoft Phi.
Relever les défis évidents
Les ingénieurs de JetBrains ont rencontré plusieurs défis critiques dans leur quête de la complétion automatique basée sur l’IA :
- Vitesse et coût : Les modèles de conversation conventionnels se sont avérés impraticables en raison de leurs coûts de calcul élevés et de leurs temps de réponse lents. Ces modèles n’étaient pas non plus conscients des techniques spécifiques au code comme le remplissage au milieu (FIM) ou la correction des jetons.
- Formatage de la sortie : Les modèles de conversation phares généraient souvent des données dans des formats incohérents, ce qui rendait difficile l’analyse des réponses et leur intégration transparente dans l’éditeur.
- Provenance des données : La détermination de l’origine des données d’entraînement et l’atténuation des problèmes potentiels de violation du droit d’auteur ont présenté un obstacle important.
Mellum : Une vue d’ensemble
L’équipe de développement de JetBrains s’est rendu compte que la création de son propre modèle était l’approche optimale. Leur objectif était de concevoir un modèle qui équilibre la qualité, les coûts d’inférence et la latence tout en étant formé sur des données avec une provenance claire. Des recherches préliminaires ont suggéré qu’un modèle avec environ 4 milliards de paramètres pourrait fournir des capacités de complétion automatique compétentes pour un large éventail de scénarios et d’utilisateurs. De plus, en entraînant exclusivement le modèle sur du code, ils pourraient établir un vocabulaire spécialisé de jetons, exempt de données non pertinentes.
Le processus d’entraînement du modèle Mellum se compose de trois étapes distinctes, chacune contribuant à de nouvelles connaissances et améliorant la qualité du code généré. La phase initiale implique un pré-entraînement de base sur un vaste corpus de fichiers individuels. La deuxième étape consiste à affiner le modèle avec un plus petit ensemble d’exemples spécialisés. L’apprentissage par renforcement avec la rétroaction de l’IA (RLAIF) est utilisé dans la troisième étape pour adapter le modèle aux caractéristiques spécifiques à l’IDE et éliminer les sorties indésirables.
Pré-entraînement
Pour éviter les problèmes potentiels liés aux origines ambiguës des données, le modèle a été entraîné à partir de zéro, nécessitant une introduction complète à de nombreux langages, à la syntaxe de programmation, aux modèles et aux concepts fondamentaux.
Ensemble de données
La principale source de données pour le pré-entraînement était TheStack. Cet ensemble de données garantit que les données sont à la fois juridiquement saines et pratiquement bénéfiques.
Processus de pré-entraînement
Pendant le pré-entraînement, l’ensemble de données a été échantillonné plusieurs fois pour atteindre environ 3 billions de jetons. Une fenêtre contextuelle de 8192 jetons a été utilisée, avec l’ensemble de données divisé en fragments de cette taille. La transformation de remplissage au milieu (FIM) a été appliquée à la moitié des fichiers de chaque fragment, encourageant le modèle à prendre en compte à la fois le code précédent et le code suivant. Cette technique imite de près les scénarios réels de génération de code.
La phase de pré-entraînement a été menée sur un cluster de seize nœuds, chacun équipé de huit GPU H100. Cette étape a pris environ 15 jours, ce qui a donné lieu au modèle Mellum-base à 4 milliards de paramètres.
Le pré-entraînement crée un modèle de complétion automatique de code à usage général avec une connaissance approfondie de nombreux langages de programmation. Cependant, à ce stade, le modèle est uniquement entraîné pour prédire le jeton suivant dans un segment de fichier sélectionné au hasard. Sans contexte supplémentaire, le modèle n’a pas conscience de la structure du code et n’a aucun mécanisme pour déterminer quand arrêter de générer du code.
L’étape de réglage fin est conçue pour répondre à ces limitations.
Réglage fin contextuel
Remplissage au milieu amélioré
Contrairement au pré-entraînement, où les fragments de code sont sélectionnés au hasard pour la prédiction, le réglage fin se concentre sur la segmentation du code d’une manière plus significative, en apprenant au modèle à extraire les fragments de code qui se produisent ‘dans la nature’.
Exemples spécialisés
En pratique, la complétion automatique du code nécessite de comprendre les fichiers environnants et des contextes plus larges, englobant éventuellement des projets entiers.
Pour le prétraitement des données, l’entreprise a lancé un projet interne nommé Code Engine : un SDK multiplateforme et un ensemble d’utilitaires de console développés pour construire le contexte directement à partir de fichiers ordinaires sans nécessiter une indexation complète du projet. Ce SDK a été déployé sur un cluster MapReduce interne et utilisé pour traiter des milliers de référentiels publics, générant de nombreux exemples utiles pour l’entraînement dans un délai raisonnable.
Trouver les bons algorithmes a nécessité quelques essais et erreurs.
Réglage pour des langages spécifiques
Les petits modèles peuvent grandement bénéficier d’une spécialisation pour des langages spécifiques. Bien que le modèle de base soit entraîné sur plus de 80 langages, la plupart des utilisateurs travaillent généralement avec un ou deux seulement. Pour résoudre ce problème, JetBrains a créé plusieurs modèles spécialisés :
- mellum-all : Prend en charge la plupart des langages et dialectes disponibles dans les IDE JetBrains, mais la qualité de la complétion automatique est inférieure à celle des modèles spécialisés.
- mellum-python : Se spécialise dans Python et Jupyter.
- mellum-kotlin : Se spécialise dans Java et Kotlin.
- mellum-web : Se spécialise dans les technologies web.
L’étape finale : RLAIF
Enfin, les cas où les objectifs d’apprentissage ne correspondent pas aux attentes des utilisateurs doivent être résolus. Une phase d’entraînement supplémentaire, RLAIF — Apprentissage par renforcement avec la rétroaction de l’IA, est utilisée pour résoudre ces problèmes.
Le modèle apprend des interactions des utilisateurs et comprend comment mieux refléter les préférences des utilisateurs.
Cette approche améliore non seulement le score de qualité global, mais réduit également le nombre d’artefacts de génération ennuyeux.
Quelle est la qualité de Mellum ?
Le modèle fonctionne exceptionnellement bien pour sa taille. Voici comment il a été évalué :
- Tout d’abord, le modèle a été évalué sur un benchmark interne nommé ‘JetBrains BigCode’.
- Ensuite, il a été testé sur des benchmarks publics bien connus tels que SAFIM.
- Enfin, les statistiques d’utilisation des fonctionnalités ont été collectées et les mesures d’utilisateur ont été calculées.
Évaluation hors ligne
La collecte de données est une tâche complexe, mais la création d’une bonne métrique qui compare la suggestion originale avec la nouvelle proposée par le réseau neuronal est encore plus difficile. Nous avons mené une petite étude et nous nous sommes finalement arrêtés sur une combinaison de deux métriques principales :
EM :
- Exact Match est une idée très populaire.
- Une prédiction est considérée comme bonne si la première ligne de la complétion correspond à la première ligne de l’original, avec un prétraitement minimal.
KK :
- La métrique a été nommée d’après ses auteurs.
- Le nombre de lignes proposées de l’original divisé par le nombre de lignes dans la complétion proposée.
JetBrains BigCode
Le modèle a été évalué par rapport à un ensemble de données de référence obtenu à l’aide de l’outil interne JetBrains BigCode.
En conservant un contrôle total sur notre ensemble de données plutôt que de nous fier à des benchmarks publics, il devient possible d’évaluer de manière fiable la qualité du modèle pour divers styles et pratiques de codage.
Les résultats de notre évaluation JetBrains BigCode montrent une qualité comparable à celle des modèles populaires, mais Mellum est plus petit et plus efficace.
Qualité des suggestions sur une seule ligne (métrique EM)
Benchmarks publics
Le modèle a été évalué non seulement sur l’ensemble de données interne, mais aussi sur divers benchmarks publics, tels que le benchmark multilingue SAFIM (syntax-aware fill in the middle).
Évaluation en ligne
La métrique principale est appelée le ratio de code complété (RoCC). Il est défini comme le ratio de caractères de code écrits à l’aide de la complétion automatique de code par rapport à la quantité totale de code dans l’éditeur.
Une autre métrique importante est le taux d’acceptation (AR), qui est calculé comme le nombre de suggestions acceptées divisé par le nombre de toutes les suggestions affichées.
Ce fut un voyage complexe, mais les spécialistes de JetBrains l’ont achevé avec dignité. Au final, un modèle général et plusieurs modèles spécialisés ont été obtenus, qui sont disponibles via la plateforme JetBrains AI. Ils fonctionnent maintenant avec succès dans JetBrains AI Assistant.
Quelle est la prochaine étape ?
- Les ingénieurs de JetBrains travaillent actuellement sur un modèle pour les langages de développement web. Il pourrait être rendu public dans un avenir proche.
- Il est prévu d’augmenter simultanément le nombre de paramètres et la diversité des données. Il existe de nombreuses tâches différentes dans le codage — Mellum sera en mesure de les effectuer également. Les performances du service restent une métrique clé, de sorte que l’expansion du modèle se fera dans des limites raisonnables.