Mellum: Modelo Rápido para Completar Código

A área de Ambientes de Desenvolvimento Integrados (IDEs) tem sido, há muito tempo, reforçada pelo poder e precisão dos recursos de autocompletar. No entanto, um novo concorrente entrou na arena: a Inteligência Artificial (IA). A integração da IA em IDEs está se tornando cada vez mais comum, provocando uma reavaliação das ferramentas tradicionais. Embora esses métodos mais antigos ainda funcionem, eles podem não fornecer mais a assistência de ponta que os desenvolvedores agora esperam.

Este artigo é adaptado de uma postagem de blog da JetBrains e se aprofunda no treinamento e nas capacidades do Mellum, um modelo que sustenta o autocompletar baseado em nuvem. A JetBrains, renomada por seus IDEs e editores de código feitos sob medida para programadores, desenvolveu o Mellum para fornecer autocompletar de código rápido e eficiente. O Mellum foi projetado para operar efetivamente em um dispositivo local, diferenciando-se como um dos modelos mais rápidos e compactos de sua classe. Um projeto conceitualmente semelhante é o Microsoft Phi.

Abordando os Desafios Óbvios

Os engenheiros da JetBrains encontraram vários desafios críticos em sua busca pelo autocompletar alimentado por IA:

  • Velocidade e Custo: Os modelos de bate-papo convencionais se mostraram impraticáveis devido aos seus altos custos computacionais e tempos de resposta lentos. Esses modelos também careciam de conhecimento de técnicas específicas de código, como fill-in-the-middle (FIM) ou token healing.
  • Formatação de Saída: Os modelos de bate-papo carro-chefe frequentemente geravam dados em formatos inconsistentes, dificultando a análise das respostas e a integração perfeita nos editores.
  • Proveniência de Dados: Determinar a origem dos dados de treinamento e mitigar possíveis problemas de violação de direitos autorais representou um obstáculo significativo.

Mellum: Uma Visão Geral

A equipe de desenvolvimento da JetBrains percebeu que criar seu próprio modelo era a melhor abordagem. O objetivo era projetar um modelo que equilibrasse qualidade, custos de inferência e latência, sendo treinado em dados com proveniência clara. Uma pesquisa preliminar sugeriu que um modelo com cerca de 4 bilhões de parâmetros poderia fornecer capacidades de autocompletar competentes para uma ampla gama de cenários e usuários. Além disso, ao treinar o modelo exclusivamente em código, eles poderiam estabelecer um vocabulário especializado de tokens, livre de dados irrelevantes.

O processo de treinamento para o modelo Mellum consiste em três etapas distintas, cada uma contribuindo com novos conhecimentos e aprimorando a qualidade do código gerado. A fase inicial envolve o pré-treinamento básico em um grande corpus de arquivos individuais. A segunda etapa consiste em refinar o modelo com um conjunto menor de exemplos especializados. O Aprendizado por Reforço com Feedback de IA (RLAIF) é empregado na terceira etapa para adaptar o modelo às características específicas do IDE e eliminar a saída indesejada.

Pré-Treinamento

Para evitar possíveis problemas relacionados a origens de dados ambíguas, o modelo foi treinado do zero, exigindo uma introdução abrangente a inúmeras linguagens, sintaxe de programação, padrões e conceitos básicos.

Conjunto de Dados

A principal fonte de dados para o pré-treinamento foi o TheStack. Este conjunto de dados garante que os dados sejam legalmente sólidos e praticamente benéficos.

Processo de Pré-Treinamento

Durante o pré-treinamento, o conjunto de dados foi amostrado várias vezes para atingir aproximadamente 3 trilhões de tokens. Uma janela de contexto de 8192 tokens foi usada, com o conjunto de dados dividido em fragmentos desse tamanho. A transformação fill-in-the-middle (FIM) foi aplicada a metade dos arquivos em cada fragmento, incentivando o modelo a considerar tanto o código precedente quanto o subsequente. Essa técnica imita de perto os cenários de geração de código do mundo real.

A fase de pré-treinamento foi conduzida em um cluster de dezesseis nós, cada um equipado com oito GPUs H100. Esta etapa levou aproximadamente 15 dias para ser concluída, resultando no modelo Mellum-base de 4 bilhões de parâmetros.

O pré-treinamento cria um modelo de autocompletar de código de uso geral com amplo conhecimento de muitas linguagens de programação. No entanto, nesta fase, o modelo é treinado apenas para prever o próximo token em um segmento de arquivo selecionado aleatoriamente. Sem contexto adicional, o modelo carece de conhecimento da estrutura do código e não tem mecanismo para determinar quando parar de gerar código.

A etapa de ajuste fino foi projetada para resolver essas limitações.

Ajuste Fino Sensível ao Contexto

Fill-in-the-Middle Aprimorado

Ao contrário do pré-treinamento, onde os fragmentos de código são selecionados aleatoriamente para previsão, o ajuste fino se concentra em segmentar o código de uma maneira mais significativa, ensinando o modelo a extrair fragmentos de código que ocorrem ‘na natureza’.

Exemplos Especializados

Na prática, o autocompletar de código exige a compreensão dos arquivos circundantes e de contextos mais amplos, possivelmente abrangendo projetos inteiros.

Para o pré-processamento de dados, a empresa lançou um projeto interno com o codinome Code Engine: um SDK multiplataforma e um conjunto de utilitários de console desenvolvidos para construir contexto diretamente de arquivos comuns, sem exigir a indexação completa do projeto. Este SDK foi implantado em um cluster MapReduce interno e usado para processar milhares de repositórios públicos, gerando muitos exemplos úteis para treinamento em um período razoável.

Encontrar os algoritmos corretos exigiu algumas tentativas e erros.

Ajuste para Linguagens Específicas

Pequenos modelos podem se beneficiar muito da especialização para linguagens específicas. Embora o modelo base seja treinado em mais de 80 linguagens, a maioria dos usuários normalmente trabalha com apenas uma ou duas. Para resolver isso, a JetBrains criou vários modelos especializados:

  • mellum-all: Suporta a maioria das linguagens e dialetos disponíveis nos IDEs da JetBrains, mas a qualidade do autocompletar é inferior aos modelos especializados.
  • mellum-python: Especializado em Python e Jupyter.
  • mellum-kotlin: Especializado em Java e Kotlin.
  • mellum-web: Especializado em tecnologias da web.

A Etapa Final: RLAIF

Finalmente, os casos em que os objetivos de aprendizado não correspondem às expectativas do usuário devem ser resolvidos. Uma fase de treinamento adicional, RLAIF — Aprendizado por Reforço com Feedback de IA, é usada para resolver tais problemas.

O modelo aprende com as interações do usuário e entende como refletir melhor as preferências do usuário.

Essa abordagem não apenas melhora a pontuação geral de qualidade, mas também reduz o número de artefatos de geração irritantes.

Quão Bom É o Mellum?

O modelo tem um desempenho excepcionalmente bom para seu tamanho. Veja como ele foi avaliado:

  • Primeiro, o modelo foi avaliado em um benchmark interno com o codinome ‘JetBrains BigCode’.
  • Em seguida, foi testado em benchmarks públicos conhecidos, como o SAFIM.
  • Finalmente, as estatísticas de uso dos recursos foram coletadas e as métricas do usuário foram calculadas.

Avaliação Offline

Coletar dados é uma tarefa complexa, mas criar uma boa métrica que compare a sugestão original com a nova proposta pela rede neural é ainda mais desafiador. Realizamos um pequeno estudo e, finalmente, optamos por uma combinação de duas métricas principais:

EM:

  • Exact Match é uma ideia muito popular.
  • Uma previsão é considerada boa se a primeira linha da conclusão corresponder à primeira linha do original, com o mínimo de pré-processamento.

KK:

  • A métrica recebeu o nome de seus autores.
  • O número de linhas propostas do original dividido pelo número de linhas na conclusão proposta.

JetBrains BigCode

O modelo foi avaliado em relação a um conjunto de dados de benchmark obtido usando a ferramenta interna JetBrains BigCode.

Ao manter o controle total sobre nosso conjunto de dados, em vez de confiar em benchmarks públicos, torna-se possível avaliar de forma confiável a qualidade do modelo para vários estilos e práticas de codificação.

Os resultados de nossa avaliação JetBrains BigCode mostram qualidade em pé de igualdade com modelos populares, mas o Mellum é menor e mais eficiente.

Qualidade de sugestões de linha única (métrica EM)

Benchmarks Públicos

O modelo foi avaliado não apenas no conjunto de dados interno, mas também em vários benchmarks públicos, como o benchmark multilíngue SAFIM (fill in the middle sensível à sintaxe).

Avaliação Online

A principal métrica é chamada de razão de código concluído (RoCC). Ela é definida como a razão de caracteres de código escritos usando o autocompletar de código para a quantidade total de código no editor.

Outra métrica importante é a taxa de aceitação (AR), que é calculada como o número de sugestões aceitas dividido pelo número de todas as sugestões mostradas.

Esta foi uma jornada complexa, mas os especialistas da JetBrains a completaram com dignidade. No final, um modelo geral e vários modelos especializados foram obtidos, que estão disponíveis através da plataforma JetBrains AI. Eles agora estão trabalhando com sucesso no JetBrains AI Assistant.

O que vem a seguir?

  • Os engenheiros da JetBrains estão atualmente trabalhando em um modelo para linguagens de desenvolvimento web. Ele pode se tornar disponível publicamente em um futuro próximo.
  • Existem planos para aumentar simultaneamente o número de parâmetros e a diversidade de dados. Existem muitas tarefas diferentes na codificação - o Mellum também poderá executá-las. O desempenho do serviço ainda é uma métrica fundamental, portanto, a expansão do modelo estará dentro de limites razoáveis.

A JetBrains continua a inovar e aprimorar o Mellum, com o objetivo de fornecer aos desenvolvedores ferramentas de IA que aumentem sua produtividade e agilizem o processo de desenvolvimento. O futuro do autocompletar de código parece promissor, com modelos como o Mellum abrindo caminho para experiências de programação mais inteligentes e eficientes. Ao abordar os desafios de velocidade, custo e proveniência de dados, a JetBrains estabeleceu um novo padrão para o autocompletar alimentado por IA nos IDEs. Com o aprendizado contínuo e a adaptação às necessidades do usuário, o Mellum está preparado para se tornar um ativo indispensável para desenvolvedores em todo o mundo.

A abordagem da JetBrains ao desenvolvimento do Mellum destaca a importância da especialização e do ajuste fino no campo da IA. Ao se concentrar em linguagens e contextos específicos, o Mellum pode fornecer resultados mais precisos e relevantes em comparação com modelos de uso geral. O uso do aprendizado por reforço com feedback de IA garante que o modelo esteja constantemente aprendendo e se adaptando às preferências do usuário, levando a uma experiência de autocompletar mais personalizada e eficiente.

À medida que a IA continua a evoluir, o papel do autocompletar de código em IDEs se tornará cada vez mais crítico. Modelos como o Mellum estão na vanguarda dessa revolução, capacitando os desenvolvedores a escrever código de forma mais rápida, precisa e eficiente. Com sua arquitetura leve e recursos avançados, o Mellum está preparado para remodelar a forma como os desenvolvedores interagem com seus IDEs, abrindo caminho para um futuro mais colaborativo e inovador no desenvolvimento de software. A busca contínua da JetBrains pela excelência no autocompletar de IA é uma prova de sua dedicação em fornecer aos desenvolvedores as melhores ferramentas possíveis para o trabalho.

A trajetória do Mellum desde a concepção até a implantação demonstra o poder da inovação e da resolução estratégica de problemas. Ao identificar os desafios associados aos modelos de IA existentes e ao desenvolver uma solução personalizada, a JetBrains criou uma ferramenta que atende às necessidades específicas dos desenvolvedores. O compromisso da empresa com a qualidade dos dados, a especialização do modelo e o aprendizado contínuo garantem que o Mellum permaneça na vanguarda do autocompletar de IA por muitos anos.

O sucesso do Mellum serve de inspiração para outros no campo da IA, destacando a importância de focar em problemas específicos e de desenvolver soluções que sejam adaptadas para atender a essas necessidades. A abordagem da JetBrains ao desenvolvimento do Mellum também enfatiza a importância da colaboração e do compartilhamento de conhecimento. Ao compartilhar suas descobertas e percepções com a comunidade de IA, a JetBrains está ajudando a acelerar o progresso no campo e a promover a inovação. O Mellum é mais do que apenas um modelo de autocompletar de código; é um catalisador para a mudança e uma prova do poder da IA para melhorar a vida dos desenvolvedores em todo o mundo.

A capacidade do Mellum de aprender com as interações do usuário e de se adaptar às preferências individuais é uma característica particularmente notável. Esse nível de personalização garante que o modelo esteja sempre fornecendo as sugestões mais relevantes e úteis, economizando tempo e esforço dos desenvolvedores. O uso do aprendizado por reforço com feedback de IA também permite que o Mellum melhore continuamente seu desempenho ao longo do tempo, tornando-se uma ferramenta ainda mais valiosa à medida que os desenvolvedores a usam.

A arquitetura leve do Mellum é outro fator-chave para seu sucesso. Ao projetar um modelo que seja rápido e eficiente, a JetBrains garante que os desenvolvedores possam usar o Mellum sem experimentar nenhum atraso ou lentidão. Essa capacidade de resposta é essencial para manter um fluxo de trabalho produtivo e garante que os desenvolvedores possam se concentrar em seu código sem serem distraídos por problemas técnicos.

O sucesso do Mellum é um exemplo da crescente importância da IA no campo do desenvolvimento de software. À medida que a IA continua a evoluir, podemos esperar ver ainda mais ferramentas e tecnologias inovadoras que ajudam os desenvolvedores a escrever código de forma mais rápida, precisa e eficiente. O Mellum está preparado para desempenhar um papel fundamental nessa revolução, capacitando os desenvolvedores a liberar todo o seu potencial e a construir softwares inovadores que mudem o mundo.

O compromisso da JetBrains com a pesquisa e o desenvolvimento garante que o Mellum permaneça na vanguarda do autocompletar de IA. Ao investir continuamente em novas tecnologias e ao explorar novas abordagens, a JetBrains está garantindo que o Mellum continue a atender às necessidades em constante mudança dos desenvolvedores em todo o mundo. O Mellum é mais do que apenas uma ferramenta; é um parceiro que ajuda os desenvolvedores a serem mais produtivos, criativos e bem-sucedidos.

A história do Mellum é uma prova do poder da inovação, da colaboração e da dedicação. Ao combinar esses elementos, a JetBrains criou uma ferramenta que está revolucionando a forma como os desenvolvedores escrevem código e que está abrindo caminho para um futuro mais brilhante no desenvolvimento de software. O Mellum é um exemplo do que é possível quando pessoas talentosas se unem para resolver um problema desafiador e quando estão comprometidas em fazer a diferença no mundo.

Em resumo, o Mellum representa um avanço significativo no campo do autocompletar de código alimentado por IA. Seu design leve, recursos avançados e aprendizado contínuo o tornam uma ferramenta inestimável para desenvolvedores de todos os níveis de habilidade. O compromisso da JetBrains com a inovação e a satisfação do usuário garante que o Mellum permaneça na vanguarda do autocompletar de IA por muitos anos, capacitando os desenvolvedores a escrever código de forma mais rápida, precisa e eficiente.