Ensinando a IA a Corrigir Seu Código: Meu Verão Melhorando o Quick Fix na Databricks

Como humanos, aprendemos a fazer coisas novas, como balé ou boxe (ambas as atividades que tive a oportunidade de experimentar neste verão!), através de tentativa e erro. Melhoramos tentando coisas novas, aprendendo com nossos erros e ouvindo orientações. Conheço bem esse ciclo de feedback - parte do meu projeto de estágio para o verão foi ensinar um modelo de recompensa a identificar melhores correções de código para mostrar aos usuários, como parte do esforço da Databricks para construir um Assistente de Código de primeira linha.

No entanto, o meu modelo não foi o único a aprender através de tentativa e erro. Ao ensinar meu modelo a distinguir boas correções de código das ruins, aprendi a escrever um código robusto, equilibrar preocupações de latência e qualidade para um produto impactante, comunicar claramente a uma equipe maior e, acima de tudo, me divertir ao longo do caminho.

Assistente Quick Fix da Databricks

Se você já escreveu código e tentou executá-lo, apenas para encontrar um erro irritante, então você apreciaria o Quick Fix. Integrado aos Notebooks e Editores SQL da Databricks, o Quick Fix é projetado para correções de alta confiança que podem ser geradas em 1-3 segundos - ideal para erros de sintaxe, nomes de colunas mal escritos e erros simples de tempo de execução. Quando o Quick Fix é acionado, ele pega o código e uma mensagem de erro, e então usa um LLM para gerar uma correção direcionada para resolver o erro.

Qual problema o meu projeto de estágio abordou?

Embora o Quick Fix já existisse e estivesse ajudando os usuários do Databricks a corrigir seu código, havia muitas maneiras de torná-lo ainda melhor! Por exemplo, depois de gerarmos uma correção de código e realizarmos algumas verificações básicas para garantir que ela segue as convenções de sintaxe, como podemos garantir que a correção que acabamos mostrando ao usuário é a mais relevante e precisa? Entre na amostragem do melhor de k - gere várias sugestões de correção possíveis e, em seguida, use um modelo de recompensa para escolher a melhor.

Estrutura do meu projeto

Meu projeto envolveu uma mistura de implementação de backend e experimentação de pesquisa, que achei divertida e cheia de aprendizado.

Gerando várias sugestões

Primeiro, expandi o fluxo de backend de Correção Rápida para gerar sugestões diversas em paralelo usando diferentes prompts e contextos. Experimentei técnicas como adicionar raciocínio em cadeia de pensamento, raciocínio de saídas previstas, variações de prompts do sistema e contexto de banco de dados seletivo para maximizar a qualidade e a diversidade das sugestões. Descobrimos que gerar sugestões com raciocínio adicional aumentou nossas métricas de qualidade, mas também induziu algum custo de latência.

Escolhendo a melhor sugestão de correção para mostrar ao usuário

Depois que várias sugestões são geradas, temos que escolher a melhor para retornar. Comecei implementando uma linha de base de votação majoritária simples, que apresentava ao usuário a correção mais frequentemente sugerida - operando no princípio de que uma solução gerada mais comumente provavelmente seria a mais eficaz. Essa linha de base teve um bom desempenho nas avaliações offline, mas não teve um desempenho significativamente melhor do que a implementação atual nos testes A/B de usuários online, então não foi implementada em produção.

Além disso, desenvolvi modelos de recompensa para classificar e selecionar as sugestões mais promissoras. Treinei os modelos para prever quais correções os usuários aceitariam e executariam com sucesso. Usamos abordagens clássicas de aprendizado de máquina (regressão logística e árvore de decisão impulsionada por gradiente usando o pacote LightGBM) e ajustamos os LLMs.

Resultados e impacto

Surpreendentemente, para a tarefa de prever a aceitação do usuário e o sucesso da execução das correções candidatas, os modelos clássicos se saíram comparativamente bem aos LLMs ajustados em avaliações offline. O modelo de árvore de decisão, em particular, pode ter se saído bem porque as edições de código que “parecem certas” para os tipos de erros que a Correção Rápida lida tendem a ser corretas: os recursos que se mostraram particularmente informativos foram a similaridade entre a linha original de código e a correção gerada, bem como o tipo de erro.

Dado esse desempenho, decidimos implantar o modelo de árvore de decisão (LightGBM) em produção. Outro fator a favor do modelo LightGBM foi seu tempo de inferência significativamente mais rápido em comparação com o LLM ajustado. A velocidade é crítica para o Quick Fix, pois as sugestões devem aparecer antes que o usuário edite manualmente seu código, e qualquer latência adicional significa menos erros corrigidos. O pequeno tamanho do modelo LightGBM tornou-o muito mais eficiente em termos de recursos e mais fácil de produzir - juntamente com algumas otimizações de modelo e infraestrutura, conseguimos diminuir nosso tempo médio de inferência em quase 100x.

Com a abordagem de melhor entre k e o modelo de recompensa implementados, conseguimos aumentar nossa taxa de aceitação interna, aumentando a qualidade para nossos usuários. Também conseguimos manter nossa latência dentro dos limites aceitáveis de nossa implementação original.

Se você quiser saber mais sobre o Assistente Databricks, confira a página de destino ou o Anúncio do Assistente Quick Fix.

Minha Experiência de Estágio

Cultura Databricks em ação

Este estágio foi uma experiência incrível para contribuir diretamente para um produto de alto impacto. Obtive uma visão direta de como a cultura da Databricks incentiva uma forte tendência para ação enquanto mantém um alto padrão para a qualidade do sistema e do produto.

Desde o início, notei como todos eram inteligentes e, ao mesmo tempo, humildes. Essa impressão só se fortaleceu com o tempo, à medida que vi o quão genuinamente solidária a equipe era. Mesmo engenheiros muito seniores regularmente se esforçavam para me ajudar a ter sucesso, seja discutindo desafios técnicos, oferecendo feedback ponderado ou compartilhando suas abordagens e aprendizados passados.

Gostaria de agradecer especialmente ao meu mentor Will Tipton, aos meus gerentes Phil Eichmann e Shanshan Zheng, aos meus mentores informais Rishabh Singh e Matt Hayes, à equipe Editor / Assistant, à equipe Applied AI e às pessoas da MosaicML por sua orientação. Aprendi habilidades inestimáveis e lições de vida com eles, que levarei comigo pelo resto da minha carreira.

Os outros estagiários incríveis!

Por último, mas não menos importante, me diverti muito conhecendo os outros estagiários! A equipe de recrutamento organizou muitos eventos divertidos que nos ajudaram a nos conectar - um dos meus favoritos foi a Olimpíada de Estagiários (foto abaixo). Seja conversando durante o almoço, experimentando aulas de ginástica locais ou comemorando aniversários com karaokê, eu realmente apreciei o quão solidário e unido o grupo de estagiários era, tanto dentro quanto fora do trabalho.

Olimpíadas de Estagiários! Vamos, Equipe 2!

Um agradecimento especial aos outros estagiários que tentaram boxe comigo!

Este verão me ensinou que o melhor aprendizado acontece quando você está resolvendo problemas reais com restrições reais - especialmente quando você está cercado por pessoas inteligentes, motivadas e solidárias. A parte mais gratificante do meu estágio não foi apenas completar o treinamento do modelo ou apresentar resultados interessantes para a equipe, mas perceber que cresci na minha capacidade de fazer melhores perguntas, raciocinar sobre trade-offs de design e entregar um recurso concreto do início ao fim em uma plataforma tão amplamente utilizada como a Databricks.

Se você quer trabalhar em projetos de ponta com colegas de trabalho incríveis, eu recomendaria que você se candidatasse para trabalhar na Databricks! Visite a página de Carreiras da Databricks para saber mais sobre as vagas disponíveis na empresa.

(This blog post has been translated using AI-powered tools) Original Post