Arquitetura de Busca Híbrida A Engenharia por Trás do Papers with Code

O Desafio de Buscar em Pesquisas Acadêmicas

Buscar por artigos de pesquisa em IA não é a mesma coisa que buscar por conteúdo web comum. Um mecanismo de busca robusto para pesquisas precisa lidar com consultas exatas, como um ID específico do arXiv, mas também precisa entender consultas conceituais. Por exemplo, um usuário que busca por "small language models for code generation" espera resultados mesmo que essas palavras exatas não apareçam no texto de um artigo. Ele também deve reconhecer solicitações de navegação como "the original BERT paper" e ser tolerante a erros de digitação ou títulos incompletos.

Para atender a essas demandas, a equipe por trás do Papers with Code não se apoiou em um único método. Eles construíram um sistema de busca híbrida sofisticado. Uma parte fundamental desse sistema depende de várias ferramentas e boas práticas em Python para gerenciar dados e computação.

Este post detalha a arquitetura que eles construíram, focando nas decisões de design que a tornam poderosa e rápida.

Uma História de Dois Métodos de Recuperação: Busca por Palavras-Chave e por Vetores

A ideia central é que a busca por palavras-chave e a busca por vetores são complementares. A busca por palavras-chave (como a busca full-text do PostgreSQL) é excelente para encontrar correspondências exatas, identificadores e nomes raros. Por outro lado, a busca por vetores usa embeddings densos para encontrar termos semanticamente semelhantes, capturando o significado 'difuso' por trás de uma consulta.

A Arquitetura: Offline vs. Online

O sistema é deliberadamente dividido em uma construção de corpus offline e um serviço de busca online. Essa separação é fundamental para equilibrar custo, desempenho e confiabilidade.

1. Construção de Corpus Offline (Jobs): O trabalho caro e orientado à produção de incorporar todo o corpus de artigos (mais de 110.000 artigos) é feito como um job em lote em uma GPU. Isso é potencializado pelos Hugging Face Jobs, que fornecem computação 'burstable'.
2. Armazenamento Durável (Buckets): A saída desses jobs (os embeddings de vetores) é armazenada em um Hugging Face Storage Bucket. Isso atua como o tecido conjuntivo, garantindo uma transferência durável entre a computação efêmera e o banco de dados de produção.
3. Serviço de Busca Online (Inference Endpoints): Apenas o pequeno trabalho sensível à latência de incorporar uma consulta de usuário ao vivo está no caminho da solicitação. Isso é tratado por um Hugging Face Inference Endpoint, projetado para respostas de baixa latência.

Aqui está uma visão simplificada do processo de incorporação de consultas:

# Pseudocódigo para o processo de incorporação de consultas online
import requests

# 1. Envia a consulta do usuário para o Inference Endpoint
response = requests.post(
    "YOUR_INFERENCE_ENDPOINT_URL",
    headers={"Authorization": "Bearer YOUR_TOKEN"},
    json={"inputs": query_text}
)

# 2. Processa a resposta para obter o vetor da consulta
query_vector = response.json()["embedding"]

# 3. Usa este vetor para realizar uma busca de distância de cosseno no PostgreSQL (pgvector)
# SELECT paper_id, embedding <=> CAST(:query_vector AS halfvec(256)) AS distance
# FROM paper_embeddings
# WHERE generation_id = :active_generation
# ORDER BY embedding <=> CAST(:query_vector AS halfvec(256))
# LIMIT 50;

Padrões de Design Principais

• Contrato de Embedding Estrito: Eles tratam o formato de embedding como uma API versionada. Isso inclui a revisão do modelo, dimensões de saída, formato de entrada e método de normalização. Isso evita falhas sutis quando os modelos são atualizados ou o código é refatorado.
• Jobs Retomáveis e Reproduzíveis: Ao organizar artefatos sob prefixos de execução imutáveis e usar somas de verificação, eles podem tentar novamente com segurança jobs com falha ou rastrear qualquer resultado até sua fonte de dados e revisão de modelo exatas.
• Degradação Graciosa: O sistema é projetado para partidas a frio. Se o Inference Endpoint estiver escalando ou atingir o tempo limite, o cliente de consulta cai imediatamente para a busca lexical. Isso garante que os usuários sempre obtenham resultados, mesmo que não sejam tão ricos semanticamente.

Lições Aprendidas e o Caminho a Seguir

A experiência da equipe oferece lições valiosas para qualquer pessoa que esteja construindo recursos de busca semelhantes com IA.

1. Separe a Produção da Latência: A incorporação em lote e a incorporação de consultas são problemas de infraestrutura diferentes. Otimize para custo e produção no primeiro e para disponibilidade e velocidade no segundo.
2. Torne o Armazenamento o Contrato: Use um bucket de armazenamento como um limite explícito e auditável entre sua computação e os sistemas de produção.
3. Fixe Mais do Que o Nome do Modelo: A revisão, a dimensão e o prompt afetam a qualidade final da recuperação. Armazene e valide todos eles.
4. Projete para Partidas a Frio: Se você usar escala para zero, deve ter um fallback rápido. A busca híbrida fornece isso naturalmente.
5. Vetores Menores Podem Ser um Recurso: A Aprendizagem de Representação Matryoshka (MRL) permite que você troque qualidade por velocidade e armazenamento. Em seu piloto, 256 dimensões preservaram a recuperação usando uma fração do armazenamento.
6. A Ativação Deve Ser Segura: Novas gerações de embeddings devem ser importadas e validadas ao lado da atual antes de serem ativadas atomicamente. Isso torna a reversão uma simples mudança de configuração.

Limitações e Considerações

Embora esta arquitetura seja poderosa, é importante notar que a busca híbrida nem sempre é a melhor solução. A equipe recomenda começar com a busca por palavras-chave como uma linha de base barata e só adicionar busca semântica e/ou híbrida quando você vir um aumento razoável na qualidade da recuperação. Adicionar um re-ranker também pode melhorar os resultados, mas introduz latência e sobrecarga adicionais.

Conclusão

O sistema de busca híbrida que alimenta o Papers with Code é uma aula magistral em design de sistemas pragmático. Ao separar cargas de trabalho online e offline, usar um bucket de armazenamento como contrato e projetar para degradação graciosa, eles construíram um sistema que é poderoso e confiável. O uso de Jobs, Buckets e Inference Endpoints da Hugging Face fornece um modelo para construir infraestrutura de IA escalável.

Para aqueles que procuram otimizar ainda mais seus sistemas de recuperação, explorar os recursos de modelos mais novos é um ótimo próximo passo. Você também pode ler sobre outros desenvolvimentos interessantes no ecossistema Python, como os novos recursos no próximo lançamento do Python 3.15.

Leia também

• Escalando Python para um Cluster na Nuvem com Ray: Um Tutorial Prático na AWS
• O que vem no Python 3.15? Um Mergulho Profundo nos Recursos do Alpha 5