Los Entornos de Desarrollo Integrados (IDEs) han sido potenciados durante mucho tiempo por la potencia y la precisión de las funciones de autocompletado. Sin embargo, un nuevo contendiente ha entrado en la arena: la Inteligencia Artificial (IA). La integración de la IA en los IDEs es cada vez más omnipresente, lo que provoca una reevaluación de las herramientas tradicionales. Si bien estos métodos más antiguos aún funcionan, es posible que ya no brinden la asistencia de vanguardia que los desarrolladores esperan ahora.
Este artículo está adaptado de una publicación de blog de JetBrains, y profundiza en la formación y las capacidades de Mellum, un modelo que sustenta el autocompletado basado en la nube. JetBrains, reconocido por sus IDEs y editores de código adaptados para programadores, ha desarrollado Mellum para proporcionar un autocompletado de código rápido y eficiente. Mellum está diseñado para operar eficazmente en un dispositivo local, lo que lo distingue como uno de los modelos más rápidos y compactos de su clase. Un proyecto conceptualmente similar es Microsoft Phi.
Abordando los Desafíos Obvios
Los ingenieros de JetBrains se encontraron con varios desafíos críticos en su búsqueda del autocompletado impulsado por la IA:
- Velocidad y Costo: Los modelos de chat convencionales demostraron ser poco prácticos debido a sus altos costos computacionales y sus lentos tiempos de respuesta. Estos modelos tampoco tenían conocimiento de técnicas específicas de código como fill-in-the-middle (FIM) o token healing.
- Formato de Salida: Los modelos de chat estrella a menudo generaban datos en formatos inconsistentes, lo que dificultaba el análisis de las respuestas y su integración perfecta en el editor.
- Procedencia de los Datos: Determinar el origen de los datos de entrenamiento y mitigar los posibles problemas de infracción de derechos de autor presentó un obstáculo importante.
Mellum: Una Visión General
El equipo de desarrollo de JetBrains se dio cuenta de que crear su propio modelo era el enfoque óptimo. Su objetivo era diseñar un modelo que equilibrara la calidad, los costos de inferencia y la latencia, mientras que se entrenaba con datos con una procedencia clara. La investigación preliminar sugirió que un modelo con alrededor de 4 mil millones de parámetros podría proporcionar capacidades de autocompletado competentes para una amplia gama de escenarios y usuarios. Además, al entrenar exclusivamente el modelo con código, podrían establecer un vocabulario especializado de tokens, libre de datos irrelevantes.
El proceso de entrenamiento para el modelo Mellum consta de tres etapas distintas, cada una de las cuales aporta nuevos conocimientos y mejora la calidad del código generado. La fase inicial implica un pre-entrenamiento básico en un gran corpus de archivos individuales. La segunda etapa consiste en refinar el modelo con un conjunto más pequeño de ejemplos especializados. El Aprendizaje por Refuerzo con Retroalimentación de la IA (RLAIF) se emplea en la tercera etapa para adaptar el modelo a las características específicas del IDE y eliminar la salida no deseada.
Pre-Entrenamiento
Para evitar posibles problemas relacionados con los orígenes ambiguos de los datos, el modelo se entrenó desde cero, lo que requería una introducción exhaustiva a numerosos lenguajes, sintaxis de programación, patrones y conceptos básicos.
Conjunto de Datos
La principal fuente de datos para el pre-entrenamiento fue TheStack. Este conjunto de datos garantiza que los datos sean tanto legalmente sólidos como prácticamente beneficiosos.
Proceso de Pre-Entrenamiento
Durante el pre-entrenamiento, el conjunto de datos se muestreó varias veces para alcanzar aproximadamente 3 billones de tokens. Se utilizó una ventana de contexto de 8192 tokens, con el conjunto de datos dividido en fragmentos de este tamaño. La transformación fill-in-the-middle (FIM) se aplicó a la mitad de los archivos de cada fragmento, lo que alentó al modelo a considerar tanto el código anterior como el posterior. Esta técnica imita fielmente los escenarios de generación de código del mundo real.
La fase de pre-entrenamiento se llevó a cabo en un clúster de dieciséis nodos, cada uno equipado con ocho GPUs H100. Esta etapa tardó aproximadamente 15 días en completarse, lo que resultó en el modelo Mellum-base de 4 mil millones de parámetros.
El pre-entrenamiento crea un modelo de autocompletado de código de propósito general con un amplio conocimiento de muchos lenguajes de programación. Sin embargo, en esta etapa, el modelo solo está entrenado para predecir el siguiente token en un segmento de archivo seleccionado aleatoriamente. Sin contexto adicional, el modelo carece de conciencia de la estructura del código y no tiene ningún mecanismo para determinar cuándo dejar de generar código.
La etapa de ajuste fino está diseñada para abordar estas limitaciones.
Ajuste Fino Consciente del Contexto
Fill-in-the-Middle Mejorado
A diferencia del pre-entrenamiento, donde los fragmentos de código se seleccionan aleatoriamente para la predicción, el ajuste fino se concentra en segmentar el código de una manera más significativa, enseñando al modelo a extraer fragmentos de código que ocurren ‘en la naturaleza’.
Ejemplos Especializados
En la práctica, el autocompletado de código requiere comprender los archivos circundantes y los contextos más amplios, que posiblemente abarquen proyectos completos.
Para el preprocesamiento de datos, la empresa lanzó un proyecto interno con el nombre en clave Code Engine: un SDK multiplataforma y un conjunto de utilidades de consola desarrolladas para construir el contexto directamente a partir de archivos ordinarios sin requerir la indexación completa del proyecto. Este SDK se implementó en un clúster interno de MapReduce y se utilizó para procesar miles de repositorios públicos, generando muchos ejemplos útiles para el entrenamiento en un plazo razonable.
Encontrar los algoritmos correctos requirió algunas pruebas y errores.
Ajuste para Lenguajes Específicos
Los modelos pequeños pueden beneficiarse enormemente de la especialización para lenguajes específicos. Si bien el modelo base está entrenado en más de 80 lenguajes, la mayoría de los usuarios normalmente trabajan con solo uno o dos. Para abordar esto, JetBrains creó múltiples modelos especializados:
- mellum-all: Admite la mayoría de los lenguajes y dialectos disponibles en los IDEs de JetBrains, pero la calidad del autocompletado es menor que la de los modelos especializados.
- mellum-python: Se especializa en Python y Jupyter.
- mellum-kotlin: Se especializa en Java y Kotlin.
- mellum-web: Se especializa en tecnologías web.
El Paso Final: RLAIF
Finalmente, se deben resolver los casos en que los objetivos de aprendizaje no coinciden con las expectativas del usuario. Se utiliza una fase de entrenamiento adicional, RLAIF — Aprendizaje por Refuerzo con Retroalimentación de la IA, para resolver tales problemas.
El modelo aprende de las interacciones del usuario y comprende cómo reflejar mejor las preferencias del usuario.
Este enfoque no solo mejora la puntuación de calidad general, sino que también reduce la cantidad de artefactos de generación molestos.
¿Qué Tan Bueno Es Mellum?
El modelo funciona excepcionalmente bien para su tamaño. Así es como se evaluó:
- Primero, el modelo se evaluó en un benchmark interno con el nombre en clave ‘JetBrains BigCode’.
- Luego se probó en benchmarks públicos conocidos como SAFIM.
- Finalmente, se recopilaron estadísticas de uso para las funciones y se calcularon las métricas del usuario.
Evaluación Offline
Recopilar datos es una tarea compleja, pero crear una buena métrica que compare la sugerencia original con la nueva propuesta por la red neuronal es aún más desafiante. Realizamos un pequeño estudio y finalmente nos decidimos por una combinación de dos métricas principales:
EM:
- Exact Match es una idea muy popular.
- Una predicción se considera buena si la primera línea de la finalización coincide con la primera línea del original, con un preprocesamiento mínimo.
KK:
- La métrica lleva el nombre de sus autores.
- El número de líneas propuestas del original dividido por el número de líneas en la finalización propuesta.
JetBrains BigCode
El modelo se evaluó con un conjunto de datos de benchmark obtenido utilizando la herramienta interna JetBrains BigCode.
Al mantener el control total sobre nuestro conjunto de datos en lugar de depender de benchmarks públicos, se hace posible evaluar de manera confiable la calidad del modelo para varios estilos y prácticas de codificación.
Los resultados de nuestra evaluación de JetBrains BigCode muestran una calidad a la par con los modelos populares, pero Mellum es más pequeño y más eficiente.
Calidad de las sugerencias de una sola línea (métrica EM)
Benchmarks Públicos
El modelo se evaluó no solo en el conjunto de datos interno, sino también en varios benchmarks públicos, como el benchmark multilingüe SAFIM (syntax-aware fill in the middle).
Evaluación Online
La métrica principal se llama la relación de código completado (RoCC). Se define como la relación entre los caracteres de código escritos utilizando el autocompletado de código y la cantidad total de código en el editor.
Otra métrica importante es la tasa de aceptación (AR), que se calcula como el número de sugerencias aceptadas dividido por el número de todas las sugerencias mostradas.
Este fue un viaje complejo, pero los especialistas de JetBrains lo completaron con dignidad. Al final, se obtuvieron un modelo general y varios modelos especializados, que están disponibles a través de la plataforma JetBrains AI. Ahora están funcionando con éxito en JetBrains AI Assistant.
¿Qué sigue?
- Los ingenieros de JetBrains están trabajando actualmente en un modelo para lenguajes de desarrollo web. Puede estar disponible públicamente en un futuro próximo.
- Hay planes para aumentar simultáneamente tanto el número de parámetros como la diversidad de los datos. Hay muchas tareas diferentes en la codificación — Mellum también podrá realizarlas. El rendimiento del servicio sigue siendo una métrica clave, por lo que la expansión del modelo estará dentro de límites razonables.