Pesquisando “por vibe” milhões de páginas da Wikipedia offline

• Implementa busca por embeddings em tempo real e offline dentro do navegador, sem servidor, permitindo explorar localmente de forma interativa 6 milhões de documentos da Wikipedia em inglês
• Os embeddings originais em float32 com 384 dimensões ocupam cerca de 9 GB, grandes demais para um app no navegador; por isso, o tamanho é reduzido com base no all-minilm-l6-v2 e em embeddings médios de chunks de páginas
• Aplica Product Quantization em subespaços de 8 dimensões para reduzir o dataset para cerca de 288 MB, permitindo armazenar cerca de 2 milhões de embeddings em 96 MB
• A busca calcula distâncias no domínio comprimido, sem descompactar, e torna o cálculo cerca de 4 vezes mais rápido ao corrigir o resultado da conversão com o opcode GatherElements do ONNX
• Em dispositivos modernos, a inferência do MiniLM e o cálculo de distância de 100 mil embeddings terminam em pouco tempo, permitindo que a UI em React atualize os resultados a cada 100–300 ms

---

Busca vetorial offline rodando dentro do navegador

• O objetivo é mostrar que é possível pesquisar milhões de páginas da Wikipedia em tempo real no navegador, sem uma infraestrutura complexa de busca no servidor
• O desenho central se resume a três pontos
  • A busca vetorial pode ser rápida o bastante usando apenas Product Quantization e varredura linear
  • Usar ONNX Runtime, mas corrigir no nível de opcode as ineficiências criadas pela conversão do PyTorch
  • O navegador é rápido o bastante para realizar inferência em tempo real com WASM, e o WebGPU continua sendo uma possibilidade futura
• As condições de busca usam em conjunto a similaridade vetorial dos embeddings do corpo do texto e uma condição de igualdade sobre a primeira letra do título
  • É algo próximo a um caso de consulta de banco de dados em que algumas colunas são pesquisadas por similaridade vetorial e outras por condições de igualdade
• Por causa da restrição de ser um app offline no navegador, os embeddings precisam ser pequenos, o modelo de embeddings precisa ser open source, e o tamanho do modelo e o volume de computação também precisam ser leves
  • Em 2023, um modelo leve geralmente significa menos de 100 MB

Dados da Wikipedia em inglês e tamanho dos embeddings

• O dataset é a Wikipedia em inglês, com embeddings gerados pelo modelo all-minilm-l6-v2
• São cerca de 6 milhões de documentos; as páginas são divididas em chunks e, depois, os embeddings de cada chunk são promediados para criar o embedding da página
• As páginas são colocadas primeiro das mais longas para as mais curtas
  • A intenção é, ao carregar o banco de dados progressivamente, mostrar primeiro as páginas longas, nas quais provavelmente houve mais esforço humano
• A dimensão dos embeddings é de 384 dimensões
  • Armazenados como float32, isso dá 6M * 384 * 4 = 9GB
  • Em 96 MB cabem apenas cerca de 64 mil embeddings
• A codificação float32 original é grande demais para uma busca offline baseada em navegador

Compressão de embeddings com Product Quantization

• Product Quantization não armazena diretamente valores de ponto flutuante; em vez disso, cria várias paletas e substitui os valores de entrada por índices nessas paletas
• As paletas normalmente têm até 256 valores, então cada índice pode ter no máximo 1 byte
  • Uma paleta implícita usa uma linha numérica como -127 a 127 ou 0 a 255, deslocada e escalonada
  • Uma paleta explícita armazena 256 números em sequência
  • Product Quantization usa paletização explícita
• Na configuração mínima, se cada uma das 384 dimensões for substituída por um índice de 1 byte, o dataset fica com 6M * 384 = 2.25GB
  • Em 96 MB, seria possível armazenar cerca de 256 mil embeddings
  • Esse método ainda é ineficiente demais para o objetivo
• Ao quantizar por pontos de 2 dimensões, é possível representar um embedding de 384 dimensões com 192 índices, resultando em economia de espaço de 2 vezes
• A configuração real usa unidades de pontos de 8 dimensões
  • Representa 384 dimensões com 48 índices
  • O tamanho da paleta se mantém na faixa de cerca de 384 KB
  • O tamanho do dataset é 6M * 48 = 288M
  • Em 96 MB, dá para armazenar cerca de 2 milhões de embeddings
• Em qualquer nível de quantização, é possível avaliar precision e recall

Busca direta sem descompactar

• A vantagem do Product Quantization é permitir buscar no estado comprimido, sem descompressão
• Se n dimensões forem agrupadas em n/2 pontos de 2 dimensões, é possível pré-calcular a distância entre cada ponto da paleta e o ponto correspondente de 2 dimensões da consulta
• Depois, cada embedding consulta os valores de distância pelos índices da paleta e os soma para calcular a distância total
  • Isso exige menos trabalho do que reconstruir o ponto completo de n dimensões e então calcular a distância de cada dimensão
• Na configuração real com pontos de 8 dimensões, é possível calcular distâncias com 1/8 do trabalho necessário ao lidar com embeddings não comprimidos

Uso do Arrow como formato de troca de dados

• Como o custo de movimentação e conversão de dados é alto, é melhor serializar os embeddings comprimidos em um formato pronto para uso imediato
• Arrow é um formato orientado a colunas adequado para esse uso
  • Os embeddings e os títulos das páginas são armazenados como tabelas Arrow
  • Os índices das paletas são tratados como um array bidimensional, e os títulos como um array unidimensional de strings
  • O foco pode ficar em copiar bits para o lugar certo, em vez de parsing ou carregamento
• Como o formato de arrays do Arrow armazena apenas dados unidimensionais, são usados dois schemas para lidar com os embeddings de 48 dimensões
  • O schema de metadados usa blocos de 100 mil linhas
  • O schema de embeddings usa blocos de 100 mil * 48 linhas
  • No momento do carregamento, os embeddings são remodelados novamente
• safetensors segue os mesmos princípios de design
• JSON serializa arrays de 48 elementos como caracteres ASCII de tamanho variável, tornando a etapa de carregamento mais complexa
• Protocol Buffers usa inteiros de largura variável base 128 para inteiros com menos de 32 bits, e esse formato de inteiro atualmente não é bem suportado por compute kernels

Parsing da Wikipedia e geração de embeddings

• Para fazer parsing da marcação da Wikipedia, é usado o mediawiki parser from hell, ou mwparserfromhell
• O dataset da Wikipedia da OLM obtém o dump de dados mais recente de uma Wikipedia em determinado idioma e faz o parsing em linhas como título e corpo
  • Esse dataset está mais próximo de código que precisa ser executado do que de dados serializados inofensivos, portanto exige confiança para execução
• O parsing é realizado em todos os núcleos de uma máquina
  • Máquinas modernas têm dezenas de núcleos, o que é adequado para sistemas com menos de 100 milhões de documentos, como a Wikipedia em inglês atual, com 6 milhões de páginas
• O modelo de embeddings usado é all-minilm-l6-v2
  • É um sentence transformer leve e com bom desempenho para inglês
  • Embeda consultas e documentos no mesmo espaço
  • É um modelo muito leve, com 22 milhões de parâmetros
• Para execução em JavaScript no navegador, também é usada uma versão quantizada em 8 bits via ONNX e transformers.js
• O modelo foi treinado com sequências de 128 tokens, então sua janela de contexto é muito menor que o tamanho médio de uma página
  • As páginas são divididas em chunks e os embeddings de cada chunk são promediados para gerar o embedding médio da página

pq.js e busca por facetas baseada em varredura linear

• Ao pesquisar milhões de documentos localmente, um esquema de indexação complexo não é necessariamente obrigatório
• O objetivo é obter os poucos itens mais próximos de um ponto específico, por exemplo as 12 menores distâncias
• Quando há um array de 10 milhões de distâncias, uma coluna de faceta com 10 milhões de valores e um valor de faceta, a filtragem é feita somando 0 se a condição corresponder, ou Infinity se não corresponder, e então executando top-k
  • Em celulares modernos, é possível encontrar os 100 melhores em menos de 10 ms
  • A implementação está em filtered-topk e em pq.js
• O cálculo de distância dos embeddings comprimidos com Product Quantization pode ser escrito em PyTorch
  • Há subspaceCount paletas, e cada paleta contém codewordCount pontos de subspaceDim dimensões
• A indexação do modelo ONNX exportado a partir do PyTorch fica em um formato um tanto estranho
  • O opcode GatherElements do ONNX realiza diretamente a operação necessária
  • Com ferramentas como ONNX-modifier, é possível adicionar e remover nós do grafo de fluxo de dados do modelo ONNX exportado
• Ao substituir várias etapas de indexação por um único opcode correto, o cálculo de distância fica cerca de 4 vezes mais rápido

Cálculo em streaming e UI interativa

• O número de embeddings usados no cálculo de distância não é fixo
  • É possível calcular em streaming as distâncias entre a consulta e uma parte dos embeddings
  • Depois de tempo suficiente, roda-se top-k com as distâncias calculadas até então para atualizar os resultados da busca
• Como todos os embeddings estão localmente, a latência até a infraestrutura de busca é, na prática, 0 ms
• A meta de responsividade da UI é mostrar resultados em 100–300 ms após uma interação
  • O MiniLM pode rodar em menos de 100 ms
  • Dependendo do dispositivo de borda, pode levar menos de 15 ms
  • O cálculo de distância de 100 mil embeddings pode rodar em cerca de 10 ms
• Ao ordenar o dataset de artigos longos para curtos, a maioria dos principais resultados finais aparece rapidamente nos resultados iniciais do streaming
  • Isso porque artigos longos muitas vezes receberam mais atenção e esforço na escrita, e costumam combinar melhor com consultas de busca do que numerosos artigos stub
• Como as distâncias são calculadas globalmente e depois são aplicados o filtro de faceta e o top-k, ao mudar valores de faceta ou a quantidade de resultados para a mesma consulta, basta executar novamente apenas a filtragem
  • Essa filtragem leva menos de 10 ms, então parece instantânea

Componentes reutilizáveis de pq.js

• Muitas funções de biblioteca do app completo de busca na Wikipedia podem ser movidas para componentes reutilizáveis de pq.js
• Atualmente, muitos shapes do ONNX estão fixados antecipadamente
• Dar suporte a diferentes níveis de quantização e diferentes dimensões de embeddings permitiria uma reutilização mais ampla

[GN⁺]

Opiniões no Hacker News

• É definitivamente interessante, mas quando tentei descrever alguns termos de filosofia e psicologia, todos os itens que eu procurava só apareceram por volta da 20ª posição
  Itens mais famosos, mas menos precisos, ficavam acima; por exemplo, não importava o que eu digitasse para definir uma modalidade específica de psicoterapia, “psychotherapy” sempre era o 1º resultado
  Por outro lado, já usei o ChatGPT para encontrar uma subárea estreita cujo nome eu não lembrava, e ele acertou todas as vezes
  A ideia de um serviço de IA que encontra o nome das coisas a partir de descrições é boa, mas não sei se limitar isso à Wikipedia ou aos títulos de artigos da Wikipedia é o caminho certo, e modelos de linguagem de grande porte de uso geral já parecem fazer isso bastante bem
  Ainda assim, como prova de conceito, e pelo fato de poder rodar localmente no navegador, é realmente muito legal

  • O objetivo era mostrar o mecanismo de banco de dados e demonstrar que, depois que o navegador baixa tudo, ele pode funcionar mesmo em um ambiente totalmente isolado
    Parece haver muitos parâmetros para ajustar. Coisas como usar apenas o primeiro parágrafo do artigo ou o texto inteiro, ou pesquisar dentro de uma faixa próxima a um determinado artigo, mas ainda não mexi nisso
    A Wikipedia é um ótimo conjunto de dados de demonstração, e eu gostaria de adicionar outros conjuntos de dados também. Por exemplo, assim como procurar “mountain” no iPhoto mostra fotos com montanhas, seria interessante pesquisar em vários conjuntos de dados com um modelo multimodal como o CLIP
  • Já usei isso no estilo de explicar uma ideia e um objetivo, e receber de volta o nome da área ou palavras-chave que valeria pesquisar
    Modelos de linguagem de grande porte parecem ser o melhor mecanismo de busca fuzzy, funcionando de uma forma bastante única, mas complementar, em relação aos mecanismos de busca tradicionais

• Gosto do conceito, mas os resultados não foram bons
  Digitei “weird looking monkey” esperando resultados como o macaco-narigudo ou o macaco-de-nariz-arrebitado-dourado, mas só apareceram artigos como “Pet monkey”, “List of individual monkeys”, “Ethnoprimatology” e “Monkey”
  Foi decepcionante, porque ao colocar a mesma consulta no Google aparece exatamente o resultado que eu esperava, e eu queria descobrir outros macacos de aparência estranha que eu não conhecia

  • Foi usado diretamente um modelo sentence-transformer pronto, com mais de um ano
    O foco da demonstração era mostrar o banco de dados de embeddings, mas os embeddings em si também têm alguma utilidade
    Não registro nenhum dado analítico sobre o que as pessoas encontram ou deixam de encontrar na página, então não estou preparado para melhorar os resultados de busca
  • Editores e diretrizes da Wikipedia geralmente não gostam de adjetivos subjetivos, e “weird looking” na consulta soa como o tipo de expressão que seria evitada em artigos da Wikipedia
    Por isso fica difícil obter bons resultados a partir desse corpus de conhecimento

• A implementação é muito boa, e é legal poder fazer isso offline. Mas a qualidade dos embeddings ainda parece insuficiente
  Uma dica que pode ajudar é embutir não o artigo inteiro da Wikipedia, mas frases mais próximas de definições, ou normalmente só a primeira frase ou o primeiro parágrafo. Não sei exatamente que parte está sendo usada agora
  Meu site OneLook também oferece desde 2003 um recurso parecido para encontrar palavras e conceitos por descrição em https://onelook.com/thesaurus/
  No começo era uma busca de dicionário reverso pura, mas ao longo dos últimos 20 anos testamos embeddings de palavras, embeddings de frases e, mais recentemente, modelos de linguagem de grande porte; hoje o GPT gera candidatos para entradas que não conseguimos responder por conta própria
  Nessa tarefa, modelos de linguagem de grande porte são muito melhores do que os métodos antigos, então minha motivação para melhorar essa parte do OneLook diminuiu um pouco. Vejo com frequência pessoas dizendo que a busca por definição reversa é o principal motivo para usarem o ChatGPT

• Vi isso um pouco tarde, mas embeddings de texto, pelo menos os usados neste artigo, em geral não são muito bons para buscar por vibe
  Em geral eles estão mais próximos de comparar palavras sobrepostas ou encontrar conteúdo parecido com a consulta
  No entanto, há um artigo recente que realmente tentou atacar esse problema: “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
  O artigo tem muitos exemplos de busca por descrições abstratas, como “arquitetos que projetam edifícios”, “empresas que fazem parte de outra empresa” e “livros que influenciaram o desenvolvimento de um gênero”
  Esses embeddings parecem dar suporte muito melhor a esse tipo de busca, então seria interessante refazer a busca offline na Wikipedia do artigo linkado usando esse novo tipo de embedding

• A página não está funcionando no meu ambiente agora; model_quantized.onnx não está sendo carregado
  Enquanto eu digitava, ele já tinha baixado 19,2 MB, a cerca de 50 KB/s; se cada visitante estiver provocando isso, talvez esteja fazendo algo terrível com a conta de largura de banda do Lee Butterman

  • Estou servindo arquivos estáticos a partir de um único t2.nano, então não sei o que vai acontecer

• O que ele faz é muito impressionante, mas a qualidade dos resultados de busca não parece boa
  Pela minha experiência, sei que é realmente difícil avaliar manualmente a qualidade dos resultados de busca. É possível estar bem perto de um resultado excelente e ainda assim retornar correspondências muito piores do que isso

  • Sim. A qualidade provavelmente não é tão boa quanto a do Similar Website Finder https://explore2.marginalia.nu/ ;)
    Acho que resultados melhores viriam com embeddings de frases mais recentes, e preciso coletar mais dados

• A tecnologia é muito impressionante, mas os resultados nem tanto
  Pesquisei por “pointy building in Paris” e apareceram Tourism in Paris, Bourse de commerce (Paris), Grands Projets of François Mitterrand, List of tallest buildings and structures in the Paris region, List of tourist attractions in Paris, Palais des congrès de Paris, Landmarks in Paris, Palais de la Bourse, Lyon, Outline of Paris, Architecture of Paris
  O prédio pontudo mais famoso de Paris não apareceu em lugar nenhum
  Para esse tipo de aplicação, embeddings de frases do documento inteiro talvez não sejam a melhor opção

  • Pelo menos 5 desses provavelmente contêm a resposta que você está procurando
  • Se você está falando da Eiffel Tower, ela não é um prédio
    Acabei de verificar o artigo, e a palavra “building” aparece 19 vezes, mas na maioria delas é verbo; depois disso, é “Chrysler Building”
    Isso considerando que não exista outro prédio pontudo famoso que eu não esteja lembrando

• Parte da mágica dos mecanismos de busca está em misturar embeddings das páginas que linkam para aquela página, ou palavras-chave tradicionais de recuperação de informação, ponderadas por cliques e pontuação de autoridade
  Sem esse sinal, muita informação útil é ignorada, e os resultados deixam de parecer mágicos
  Ainda assim, é uma demonstração impressionante e interessante

• Eu queria gostar, mas nas buscas que tentei quase não apareceram resultados relevantes
  “The wizard in The Lord of the Rings” não trouxe Gandalf nem Saruman, só livros relacionados a LOTR
  “Protagonist of Scorsese's Taxi Driver” não trouxe Travis Bickle
  “A person that plants trees for a living” por algum motivo não incluiu gardener na lista
  “Curly-haired painter on TV” não trouxe Bob Ross em lugar nenhum
  Em “Unusually shaped modern art museum in Spain”, Bilbao até aparece em 4º, mas os demais não tinham formatos incomuns
  Para “Dog shaped like a sausage”, eu esperaria que dachshund estivesse entre os primeiros resultados

  • Vale destacar que todos os resultados esperados aqui têm artigos na Wikipedia
    Se não houvesse artigos, a ausência seria menos estranha, mas na prática todos existem

• “Vibes” soa muito mais intuitivo do que “sentence embeddings”. Talvez eu também deva começar a usar essa expressão :)

  • É uma expressão que sacrifica bastante a precisão em troca de “soar intuitiva”
    O autor original não explicou por que escolheu essa palavra, e ela não combina com nenhum uso de “vibe” que eu conheça
    Fico me perguntando se “gist” não era buzzword o suficiente