Buscar millones de páginas de Wikipedia sin conexión con “vibes”

 (leebutterman.com)
1 puntos por GN⁺ 2023-09-03 | 1 comentarios | Compartir por WhatsApp
  • Implementa búsqueda de embeddings en tiempo real sin conexión dentro del navegador, sin servidor, para explorar localmente de forma interactiva 6 millones de documentos de English Wikipedia
  • Los embeddings originales float32 de 384 dimensiones pesan unos 9 GB, demasiado para una app de navegador, así que se reduce el tamaño con base en all-minilm-l6-v2 y embeddings promedio de chunks de páginas
  • Aplica Product Quantization en subespacios de 8 dimensiones para reducir el dataset a unos 288 MB, y permite guardar cerca de 2 millones de embeddings en 96 MB
  • La búsqueda calcula distancias en el dominio comprimido sin descomprimir, y al corregir el resultado de la conversión con el opcode GatherElements de ONNX, el cálculo se vuelve unas 4 veces más rápido
  • En dispositivos modernos, la inferencia de MiniLM y el cálculo de distancias sobre 100k embeddings terminan en poco tiempo, lo que permite que la UI en React actualice resultados cada 100–300 ms

Búsqueda vectorial sin conexión que corre dentro del navegador

  • El objetivo es demostrar que se puede buscar en tiempo real entre millones de páginas de Wikipedia desde el navegador sin una infraestructura compleja de búsqueda en servidores
  • El diseño central se resume en tres puntos
    • La búsqueda vectorial puede ser lo suficientemente rápida solo con Product Quantization y escaneo lineal
    • Se usa ONNX Runtime, pero las ineficiencias surgidas en la conversión desde PyTorch pueden corregirse a nivel de opcode
    • El navegador es lo bastante rápido para ejecutar inferencia en tiempo real con WASM, y WebGPU queda como una posibilidad a futuro
  • Las condiciones de búsqueda combinan similitud vectorial del embedding del contenido y una condición de igualdad sobre la primera letra del título
    • Se parece a un caso de consulta de base de datos donde algunas columnas se buscan por similitud vectorial y otras por condiciones de igualdad
  • Por la restricción de ser una app de navegador sin conexión, los embeddings deben ser pequeños, el modelo de embeddings debe ser open source, y tanto el tamaño del modelo como la carga de cómputo deben ser ligeros
    • En 2023, un modelo ligero normalmente significa menos de 100 MB

Datos de English Wikipedia y tamaño de los embeddings

  • El dataset es English Wikipedia, y se generan embeddings con el modelo all-minilm-l6-v2
  • Hay cerca de 6 millones de documentos; se divide cada página en chunks y luego se promedian los embeddings de cada chunk para crear el embedding de la página
  • Las páginas se ordenan poniendo primero las más largas
    • Esto busca mostrar primero, al cargar la base de datos de forma progresiva, las páginas largas en las que probablemente las personas invirtieron más esfuerzo
  • La dimensión del embedding es de 384 dimensiones
    • Si se guarda en float32: 6M * 384 * 4 = 9GB
    • En 96 MB solo caben aproximadamente 64k embeddings
  • La codificación float32 original es demasiado grande para una búsqueda sin conexión basada en navegador

Compresión de embeddings con Product Quantization

  • Product Quantization no guarda directamente los valores de punto flotante; crea varias paletas y reemplaza los valores de entrada por índices de paleta
  • Las paletas suelen tener 256 valores o menos, por lo que cada índice puede ocupar como máximo 1 byte
    • Una paleta implícita usa una recta numérica como -127~127 o 0~255, desplazándola y escalándola
    • Una paleta explícita guarda 256 números en orden
    • Product Quantization usa paletización explícita
  • En una configuración mínima, si cada una de las 384 dimensiones se reemplaza por un índice de 1 byte, el dataset queda en 6M * 384 = 2.25GB
    • En 96 MB se pueden guardar unos 256k embeddings
    • Este enfoque sigue siendo ineficiente frente al objetivo
  • Si se cuantiza por puntos de 2 dimensiones, un embedding de 384 dimensiones se representa con 192 índices, logrando ahorrar 2 veces el espacio
  • La configuración real usa unidades de puntos de 8 dimensiones
    • Representa 384 dimensiones con 48 índices
    • El tamaño de la paleta se mantiene en alrededor de 384 KB
    • El tamaño del dataset es 6M * 48 = 288M
    • En 96 MB se pueden guardar cerca de 2 millones de embeddings
  • En cualquier nivel de cuantización se pueden evaluar precision y recall

Buscar directamente sin descomprimir

  • La ventaja de Product Quantization es que permite buscar en estado comprimido sin descomprimir
  • Si se agrupan n dimensiones en n/2 puntos de 2 dimensiones, se puede precalcular la distancia entre cada punto de la paleta y el punto correspondiente de 2 dimensiones de la consulta
  • Luego, para cada embedding, se consulta el valor de distancia mediante el índice de paleta y se suman los valores para calcular la distancia total
    • Requiere menos trabajo que reconstruir el punto n-dimensional completo y luego calcular la distancia en cada dimensión
  • En la configuración real de puntos de 8 dimensiones, el cálculo de distancias requiere alrededor de 1/8 del trabajo respecto a manejar embeddings sin comprimir

Uso de Arrow como formato de intercambio de datos

  • Como el costo de mover y transformar datos es alto, conviene serializar los embeddings comprimidos en un formato listo para usarse
  • Arrow es un formato orientado a columnas adecuado para este uso
    • Guarda embeddings y títulos de páginas como tablas Arrow
    • Los índices de paleta se tratan como un arreglo bidimensional, y los títulos como un arreglo unidimensional de strings
    • Permite concentrarse en copiar bits al lugar correcto más que en parsear o cargar
  • El formato de arreglos de Arrow solo guarda datos unidimensionales, así que se usan dos esquemas para manejar las 48 dimensiones de embedding
    • El esquema de metadata trabaja en bloques de 100 mil filas
    • El esquema de embeddings trabaja en bloques de 100 mil * 48 filas
    • Al cargar, los embeddings se vuelven a reshape
  • safetensors sigue los mismos principios de diseño
  • JSON serializa un arreglo de 48 elementos como caracteres ASCII de longitud variable, lo que complica la etapa de carga
  • Protocol Buffers usa enteros de ancho variable base 128 para enteros de menos de 32 bits, y este formato de enteros actualmente no está bien soportado por los compute kernels

Parseo de Wikipedia y generación de embeddings

  • Para parsear el marcado de Wikipedia se usa mediawiki parser from hell, es decir, mwparserfromhell
  • El dataset de Wikipedia de OLM toma el dump de datos más reciente de la Wikipedia de un idioma específico y lo parsea en filas como título y contenido
    • Este dataset se parece más a código que hay que ejecutar que a datos serializados inocuos, por lo que requiere confianza para ejecutarlo
  • El parseo se ejecuta en todos los cores de una máquina
    • Las máquinas modernas tienen decenas de cores, por lo que son adecuadas para sistemas de menos de 100 millones de documentos, como la English Wikipedia actual de 6 millones de páginas
  • El modelo de embeddings usado es all-minilm-l6-v2
    • Es un sentence transformer ligero y de buen rendimiento para inglés
    • Embebe consultas y documentos en el mismo espacio
    • Es un modelo muy liviano, con 22 millones de parámetros
  • Para ejecutarlo en JavaScript en el navegador, también se usa una versión cuantizada a 8 bits mediante ONNX y transformers.js
  • El modelo fue entrenado con secuencias de 128 tokens, por lo que su ventana de contexto es mucho más corta que la longitud promedio de una página
    • Se divide cada página en chunks y se promedian los embeddings de cada chunk para generar un embedding promedio de página

pq.js y búsqueda facetada basada en escaneo lineal

  • Al buscar localmente entre millones de documentos, no necesariamente hace falta un sistema de indexación complejo
  • El objetivo es obtener los pocos elementos más cercanos a un punto específico, por ejemplo las 12 distancias principales
  • Cuando hay un arreglo de 10 millones de distancias, una columna facet de 10 millones de valores y un único valor facet, se filtra sumando 0 si coincide e Infinity si no coincide, y luego se busca el top-k
    • En teléfonos modernos, se puede encontrar el top 100 en menos de 10 ms
    • La implementación está en filtered-topk y pq.js
  • El cálculo de distancias de embeddings comprimidos con Product Quantization puede escribirse en PyTorch
    • Hay subspaceCount paletas, y cada paleta contiene codewordCount puntos de subspaceDim dimensiones
  • La indexación del modelo ONNX exportado desde PyTorch queda con una forma algo incómoda
    • El opcode GatherElements de ONNX realiza directamente la operación necesaria
    • Con herramientas como ONNX-modifier se pueden agregar y eliminar nodos del grafo de flujo de datos del modelo ONNX exportado
  • Al reemplazar varios pasos de indexación por un único opcode correcto, el cálculo de distancias se acelera unas 4 veces

Cálculo en streaming e interfaz interactiva

  • La cantidad de embeddings sobre los que se calculan distancias no es fija
    • Se pueden calcular en streaming las distancias entre la consulta y una parte de los embeddings
    • Cuando pasa suficiente tiempo, se ejecuta top-k con las distancias calculadas hasta ese momento y se actualizan los resultados de búsqueda
  • Como todos los embeddings están en local, la latencia hacia la infraestructura de búsqueda es, en la práctica, 0 ms
  • El objetivo de respuesta de la UI es que los resultados aparezcan entre 100 y 300 ms después de una interacción
    • MiniLM puede ejecutarse en menos de 100 ms
    • En algunos dispositivos edge puede tardar menos de 15 ms
    • El cálculo de distancias sobre 100k embeddings puede ejecutarse en unos 10 ms
  • Si el dataset se ordena desde los artículos largos hacia los cortos, la mayoría de los resultados finales principales aparecen rápidamente en los resultados iniciales del streaming
    • Esto se debe a que los artículos largos suelen haber recibido más atención y esfuerzo al ser escritos, y a menudo encajan mejor con las consultas de búsqueda que numerosos artículos stub
  • Como se calculan las distancias globalmente y luego se aplican el filtro facet y el top-k, al cambiar el valor facet o la cantidad de resultados para una misma consulta solo hace falta volver a ejecutar el filtrado
    • Este filtrado toma menos de 10 ms, por lo que se siente instantáneo

Componentes reutilizables de pq.js

  • Muchas funciones de biblioteca de la app completa de búsqueda en Wikipedia pueden trasladarse a componentes reutilizables de pq.js
  • Actualmente muchas shapes de ONNX están fijadas de antemano
  • Soportar distintos niveles de cuantización y distintas dimensiones de embeddings permitiría una reutilización más amplia

Lecturas relacionadas

1 comentarios

 
GN⁺ 2023-09-03
Opiniones de Hacker News

  • Sin duda es interesante, pero cuando probé describir algunos términos de filosofía y psicología, todos los artículos que buscaba aparecieron recién alrededor del puesto 20
    Arriba quedaban artículos más famosos pero menos precisos; por ejemplo, sin importar qué escribiera para definir una modalidad específica de psicoterapia, “psychotherapy” siempre salía en primer lugar
    En cambio, alguna vez usé ChatGPT para encontrar un subcampo acotado cuyo nombre no recordaba, y lo acertó todas las veces
    La idea de un servicio de IA que encuentre el nombre de algo a partir de una descripción es buena, pero no sé si limitarlo a Wikipedia o a los títulos de artículos de Wikipedia sea el enfoque correcto, y los modelos de lenguaje grandes de propósito general ya parecen hacerlo bastante bien
    Aun así, como prueba de concepto y por poder ejecutarse localmente en el navegador, es realmente genial

    • El objetivo era mostrar el motor de base de datos y demostrar que, una vez que el navegador descarga todo, puede funcionar incluso en un entorno completamente aislado
      Parece haber muchos parámetros por ajustar: si usar solo el primer párrafo de un documento o el documento completo, si buscar dentro de un rango cercano a un documento específico, etc.; todavía no los he afinado
      Wikipedia es un excelente dataset de demostración, y me gustaría agregar otros datasets. Por ejemplo, así como al buscar “mountain” en iPhoto aparecen fotos con montañas, sería interesante buscar en varios datasets con modelos multimodales como CLIP
    • Lo he usado de una forma en la que explico una idea y un objetivo, y recibo el nombre del campo o palabras clave que valdría la pena buscar
      Los modelos de lenguaje grandes parecen el mejor motor de búsqueda difusa, y funcionan de una manera bastante singular, aunque complementaria, respecto de los motores de búsqueda tradicionales

  • Me gusta el concepto, pero los resultados no me salieron bien
    Escribí “weird looking monkey” esperando resultados como el mono narigudo o el mono dorado de nariz chata, pero solo aparecieron artículos como “Pet monkey”, “List of individual monkeys”, “Ethnoprimatology” y “Monkey”
    Al poner la misma consulta en Google obtuve exactamente los resultados que esperaba, así que fue una lástima; quería encontrar monos de aspecto raro que no conociera

    • Se usó tal cual un modelo sentence-transformer ya existente, de más de un año
      El foco de la demo estaba en mostrar la base de datos de embeddings, pero los embeddings en sí también sirven un poco
      No guardo ningún dato analítico sobre qué encuentra o no encuentra la gente en la página, así que no estoy preparado para mejorar los resultados de búsqueda
    • Los editores y las guías de Wikipedia, en general, no son muy partidarios de los adjetivos subjetivos, y el “weird looking” de la consulta suena como una expresión que se evitaría en un artículo de Wikipedia
      Por eso se vuelve difícil obtener buenos resultados desde ese corpus de conocimiento

  • La implementación está realmente buena, y es genial poder hacer esto offline. Sin embargo, la calidad de los embeddings todavía parece insuficiente
    Un truco que podría ayudar es no embeber todo el artículo de Wikipedia, sino solo oraciones cercanas a una definición o, normalmente, la primera oración o el primer párrafo. No sé bien qué parte se está usando ahora
    Mi sitio OneLook también ofrece desde 2003 una función parecida para encontrar palabras y conceptos a partir de descripciones en https://onelook.com/thesaurus/
    Al principio era una búsqueda pura de diccionario inverso, pero durante los últimos 20 años probé embeddings de palabras, embeddings de oraciones y, recientemente, hasta modelos de lenguaje grandes; hoy GPT genera candidatos para las entradas que nuestro sistema no puede responder por sí solo
    En esta tarea, los modelos de lenguaje grandes son mucho mejores que los métodos anteriores, así que se me redujo un poco la motivación para mejorar esta parte de OneLook. Veo con frecuencia a gente que dice que la búsqueda por definición inversa es una de las principales razones por las que usa ChatGPT

  • Lo vi un poco tarde, pero los embeddings de texto, al menos los usados en este artículo, normalmente no son muy buenos para buscar por vibra
    En general se acercan más a comparar palabras superpuestas o a encontrar contenido parecido a la consulta
    Sin embargo, hay un artículo reciente que sí intenta abordar este problema: “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
    El artículo contiene muchos ejemplos de búsquedas con descripciones abstractas, como “arquitectos que diseñan edificios”, “empresas que forman parte de otra empresa” y “libros que influyeron en el desarrollo de un género”
    Esos embeddings parecen soportar mucho mejor este tipo de búsqueda, así que sería interesante rehacer la búsqueda offline en Wikipedia del artículo enlazado con este nuevo tipo de embeddings

  • Ahora mismo la página no funciona en mi entorno: model_quantized.onnx no se está cargando
    Mientras escribía, ya había descargado hasta 19.2 MB a unos 50 KB/s; si cada visitante provoca esto, quizá esté causando algo terrible en la factura de ancho de banda de Lee Butterman

    • Lo estoy sirviendo como archivos estáticos desde un solo t2.nano, así que no sé qué pasará

  • Lo que hace es muy impresionante, pero la calidad de los resultados de búsqueda no parece buena
    Por experiencia sé que es realmente difícil evaluar a mano la calidad de los resultados de búsqueda. Puede estar casi llegando a un resultado muy bueno y, aun así, devolver una coincidencia mucho peor

    • Sí. Probablemente la calidad no sea tan buena como la de Similar Website Finder https://explore2.marginalia.nu/ ;)
      Creo que con embeddings de oraciones más recientes los resultados mejorarían, y tendré que recopilar más datos

  • La tecnología es muy impresionante, pero los resultados no lo fueron
    Busqué “pointy building in Paris” y salieron 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
    El edificio puntiagudo más famoso de París no aparecía por ningún lado
    Para este tipo de aplicaciones, quizá los embeddings de oraciones de documentos completos no sean lo mejor

    • Creo que al menos 5 de esos contienen la respuesta que buscas
    • Si te refieres a la Eiffel Tower, eso no es un edificio
      Acabo de revisar el artículo y la palabra “building” aparece 19 veces, pero la mayoría son verbos, y luego está “Chrysler Building”
      Esto suponiendo que no haya otro edificio puntiagudo famoso que no se me venga a la mente

  • Parte de la magia de los motores de búsqueda está en mezclar los embeddings de las páginas que enlazan a esa página, o palabras clave de recuperación de información tradicionales, ponderadas por clics y puntajes de autoridad
    Sin esa señal, se ignora mucha información útil y los resultados no se sienten mágicos
    Aun así, es una demo impresionante e interesante

  • Quería que me gustara, pero en las búsquedas que probé casi no salieron resultados relevantes
    “The wizard in The Lord of the Rings” no muestra a Gandalf ni a Saruman, solo libros relacionados con LOTR
    “Protagonist of Scorsese's Taxi Driver” no mostró a Travis Bickle
    “A person that plants trees for a living” curiosamente no incluyó gardener en la lista
    “Curly-haired painter on TV” no mostró a Bob Ross en absoluto
    En “Unusually shaped modern art museum in Spain”, Bilbao sí aparece en cuarto lugar, pero el resto no tenía formas inusuales
    Con “Dog shaped like a sausage”, pensaría que dachshund debería estar entre los primeros resultados

    • Vale la pena señalar que todos los resultados que se esperaban aquí tienen artículo en Wikipedia
      Si no tuvieran artículo, sería menos raro que faltaran, pero en realidad todos existen

  • “Vibes” es una expresión que se siente mucho más intuitiva que “sentence embeddings”. Quizá yo también debería empezar a usarla :)

    • Es una expresión que sacrifica mucha precisión a cambio de ser “intuitiva”
      El autor del post original no explicó por qué eligió esa palabra, y no encaja con ningún uso de “vibe” que yo conozca
      Me pregunto si “gist” no sonaba lo suficientemente de moda