Rechercher des millions de pages Wikipédia hors ligne, en mode « vibe »

• Mise en œuvre d’une recherche par embeddings hors ligne en temps réel directement dans le navigateur, sans serveur, pour explorer de façon interactive 6 millions d’articles de la Wikipédia anglophone en local
• Les embeddings d’origine en float32 sur 384 dimensions pèsent environ 9 Go, ce qui ne convient pas à une app navigateur ; la taille est donc réduite à partir de all-minilm-l6-v2 et d’embeddings moyens par chunk de page
• La Product Quantization est appliquée sur des sous-espaces de 8 dimensions, ce qui ramène le dataset à environ 288 Mo, avec environ 2 millions d’embeddings dans 96 Mo
• La recherche calcule les distances directement dans le domaine compressé, sans décompression, et corrige la conversion via l’opcode GatherElements d’ONNX pour rendre le calcul environ 4 fois plus rapide
• Sur les appareils récents, l’inférence MiniLM et le calcul des distances sur 100k embeddings se terminent rapidement, ce qui permet à l’interface React d’actualiser les résultats toutes les 100 à 300 ms

---

Recherche vectorielle hors ligne dans le navigateur

• L’objectif est de montrer qu’il est possible de faire de la recherche en temps réel sur des millions de pages Wikipédia dans le navigateur, sans infrastructure serveur de recherche complexe
• La conception repose sur trois idées clés
  • la recherche vectorielle peut être suffisamment rapide avec la seule Product Quantization et un balayage linéaire
  • ONNX Runtime est utilisé, mais les inefficacités issues de la conversion depuis PyTorch peuvent être corrigées au niveau des opcodes
  • le navigateur est assez rapide, via WASM, pour exécuter l’inférence en temps réel, et WebGPU reste une piste pour la suite
• Les conditions de recherche combinent la similarité vectorielle des embeddings du contenu et une condition d’égalité sur la première lettre du titre
  • cela se rapproche d’une requête de base de données où certaines colonnes sont recherchées par similarité vectorielle et d’autres par égalité
• Du fait de la contrainte d’une application navigateur hors ligne, les embeddings doivent être petits, le modèle d’embedding doit être open source, et la taille comme le coût de calcul du modèle doivent rester modestes
  • en 2023, un modèle léger signifie généralement moins de 100 Mo

Données de Wikipédia anglophone et taille des embeddings

• Le dataset est la Wikipédia anglophone, avec des embeddings produits par le modèle all-minilm-l6-v2
• Il contient environ 6 millions de documents ; chaque page est découpée en chunks, puis les embeddings de ces chunks sont moyennés pour produire un embedding de page
• Les pages sont ordonnées des plus longues aux plus courtes
  • lors du chargement progressif de la base, cela permet d’afficher d’abord les longues pages, probablement plus travaillées par des humains
• La dimension des embeddings est de 384
  • stockés en float32, cela donne 6M * 384 * 4 = 9GB
  • 96 Mo ne permettent de stocker qu’environ 64k embeddings
• L’encodage original en float32 est trop volumineux pour une recherche hors ligne dans le navigateur

Compression des embeddings avec Product Quantization

• La Product Quantization ne stocke pas directement les valeurs en virgule flottante ; elle construit plusieurs palettes puis remplace les valeurs d’entrée par des indices dans ces palettes
• Les palettes contiennent généralement au plus 256 valeurs, donc chaque indice peut tenir sur 1 octet au maximum
  • une palette implicite utilise une droite numérique comme -127 à 127 ou 0 à 255, déplacée et redimensionnée
  • une palette explicite stocke les 256 nombres dans l’ordre
  • la Product Quantization utilise une quantification par palette explicite
• Dans la configuration minimale, si chaque dimension parmi les 384 est remplacée par un indice sur 1 octet, le dataset devient 6M * 384 = 2.25GB
  • 96 Mo peuvent alors contenir environ 256k embeddings
  • cette approche reste encore trop inefficace pour l’objectif visé
• En quantifiant par points 2D, un embedding de 384 dimensions peut être représenté par 192 indices, soit un gain d’espace par 2
• La configuration réelle utilise des points de 8 dimensions
  • les 384 dimensions sont représentées par 48 indices
  • la taille des palettes reste autour de 384 Ko
  • la taille du dataset devient 6M * 48 = 288M
  • 96 Mo peuvent contenir environ 2 millions d’embeddings
• Il est possible d’évaluer precision et recall pour chaque niveau de quantification

Recherche directe sans décompression

• L’intérêt de la Product Quantization est qu’elle permet de faire la recherche directement sur les données compressées, sans décompression
• Si l’on regroupe n dimensions en n/2 points 2D, on peut précalculer la distance entre chaque point de palette et le point 2D correspondant de la requête
• Ensuite, pour chaque embedding, il suffit de lire les valeurs de distance via les indices de palette et de les additionner pour obtenir la distance totale
  • cela demande moins de travail que de reconstruire d’abord tout le point n-dimensionnel puis de calculer les distances dimension par dimension
• Dans la configuration réelle à 8 dimensions par point, le calcul de distance demande environ 1/8 du travail nécessaire avec des embeddings non compressés

Utiliser Arrow comme format d’échange de données

• Comme le coût de déplacement et de transformation des données est élevé, il est avantageux de sérialiser les embeddings compressés dans un format directement exploitable
• Arrow est un format en colonnes adapté à cet usage
  • les embeddings et les titres de pages sont stockés dans des tables Arrow
  • les indices de palette sont traités comme un tableau 2D, les titres comme un tableau 1D de chaînes
  • l’objectif est de se concentrer sur la copie des bits au bon endroit plutôt que sur le parsing ou le chargement
• Comme le format de tableau Arrow ne stocke que des données 1D, deux schémas sont utilisés pour manipuler les 48 dimensions des embeddings
  • le schéma de métadonnées fonctionne par blocs de 100 000 lignes
  • le schéma des embeddings fonctionne par blocs de 100000 * 48 lignes
  • au chargement, les embeddings sont de nouveau remis en forme
• safetensors suit le même principe de conception
• JSON sérialise les tableaux de 48 éléments en caractères ASCII de longueur variable, ce qui complique le chargement
• Protocol Buffers utilise des entiers à largeur variable en base 128 pour les entiers inférieurs à 32 bits, un format actuellement mal pris en charge par les compute kernels

Parsing de Wikipédia et génération des embeddings

• Pour parser le balisage Wikipédia, le projet utilise mediawiki parser from hell, c’est-à-dire mwparserfromhell
• Le dataset Wikipédia d’OLM récupère le dump de données le plus récent d’une Wikipédia donnée, puis le parse en lignes comme le titre et le corps du texte
  • ce dataset ressemble moins à des données sérialisées inoffensives qu’à du code qu’il faut exécuter, ce qui suppose de lui faire confiance
• Le parsing s’exécute sur tous les cœurs d’une seule machine
  • les machines récentes disposent de dizaines de cœurs, ce qui convient à des ensembles de moins de 100 millions de documents comme la Wikipédia anglophone actuelle et ses 6 millions de pages
• Le modèle d’embedding utilisé est all-minilm-l6-v2
  • c’est un sentence transformer léger et performant pour l’anglais
  • il encode les requêtes et les documents dans le même espace
  • avec 22 millions de paramètres, c’est un modèle très léger
• Une version quantifiée en 8 bits est aussi utilisée via ONNX et transformers.js pour l’exécution JavaScript dans le navigateur
• Le modèle a été entraîné sur des séquences de 128 tokens, bien plus courtes que la longueur moyenne d’une page
  • les pages sont donc découpées en chunks, puis leurs embeddings sont moyennés afin de produire un embedding moyen de page

Recherche par facettes avec pq.js et balayage linéaire

• Quand on recherche localement dans plusieurs millions de documents, une structure d’indexation complexe n’est pas forcément nécessaire
• Le but est d’obtenir les quelques points les plus proches d’un point donné, par exemple les 12 plus petites distances
• Avec un tableau de 10 millions de distances, 10 millions de valeurs de facette et une valeur de facette cible, le filtrage consiste à ajouter 0 si la condition est satisfaite, sinon Infinity, puis à chercher le top-k
  • sur un smartphone récent, on peut trouver les 100 premiers résultats en moins de 10 ms
  • l’implémentation se trouve dans filtered-topk et pq.js
• Le calcul des distances sur des embeddings compressés par Product Quantization peut être écrit en PyTorch
  • il y a subspaceCount palettes, et chaque palette contient codewordCount points de dimension subspaceDim
• L’indexation du modèle ONNX exporté depuis PyTorch prend une forme un peu maladroite
  • l’opcode GatherElements d’ONNX exécute directement l’opération nécessaire
  • avec des outils comme ONNX-modifier, il est possible d’ajouter ou supprimer des nœuds dans le graphe de flux de données d’un modèle ONNX exporté
• En remplaçant plusieurs étapes d’indexation par un seul opcode correct, le calcul des distances devient environ 4 fois plus rapide

Calcul en streaming et interface interactive

• Le nombre d’embeddings sur lesquels calculer les distances n’est pas fixe
  • il est possible de calculer en streaming les distances entre la requête et une partie des embeddings
  • après un certain temps, on peut exécuter le top-k sur les distances déjà calculées pour mettre à jour les résultats de recherche
• Comme tous les embeddings sont stockés en local, la latence jusqu’à l’infrastructure de recherche est en pratique de 0 ms
• L’objectif de réactivité de l’interface est d’afficher des résultats 100 à 300 ms après une action
  • MiniLM peut s’exécuter en moins de 100 ms
  • sur certains appareils edge, cela peut descendre sous les 15 ms
  • le calcul des distances sur 100k embeddings peut se faire en environ 10 ms
• Si le dataset est trié des articles les plus longs aux plus courts, la plupart des meilleurs résultats finaux apparaissent rapidement dans les premiers résultats du streaming
  • les articles longs demandent souvent plus d’attention et d’effort à la rédaction, et correspondent souvent mieux à la requête que les très nombreux articles très courts
• Comme les distances sont calculées globalement avant d’appliquer le filtre par facette et le top-k, il suffit de relancer le filtrage quand on change la valeur de facette ou le nombre de résultats pour une même requête
  • ce filtrage prend moins de 10 ms et paraît instantané

Composants pq.js réutilisables

• Une grande partie des fonctions de bibliothèque de l’ensemble de l’app Wikipedia search app peut être déplacée vers des composants réutilisables de pq.js
• À l’heure actuelle, beaucoup de shapes ONNX sont figées à l’avance
• Une prise en charge de différents niveaux de quantification et de différentes dimensions d’embedding permettrait une réutilisation plus large

[GN⁺]

Commentaires sur Hacker News

• C’est clairement intéressant, mais quand j’ai essayé de décrire quelques termes de philosophie et de psychologie, les entrées que je cherchais n’apparaissaient toutes qu’autour de la 20e place
  Des entrées plus connues mais moins exactes étaient placées au-dessus ; par exemple, quoi que je saisisse pour définir une modalité de psychothérapie précise, “psychotherapy” était toujours en première position
  À l’inverse, il m’est déjà arrivé d’utiliser ChatGPT pour retrouver un sous-domaine étroit dont je ne me rappelais plus le nom, et il a trouvé juste à chaque fois
  L’idée d’un service d’IA qui retrouve le nom des choses à partir d’une description est bonne, mais je ne sais pas si la limiter à Wikipedia ou aux titres d’articles Wikipedia est la bonne approche, et les grands modèles de langage généralistes semblent déjà plutôt bien s’en sortir
  Cela dit, comme preuve de concept, et parce que ça peut tourner localement dans le navigateur, c’est vraiment chouette

  • L’objectif était de montrer le moteur de base de données et de démontrer qu’une fois que le navigateur a tout téléchargé, il peut fonctionner même dans un environnement totalement isolé
    Il semble y avoir beaucoup de paramètres à ajuster : utiliser seulement le premier paragraphe d’un article ou l’article entier, chercher dans un périmètre proche d’un article donné, etc., mais je n’y ai pas encore touché
    Wikipedia est un excellent jeu de données de démonstration, et j’aimerais aussi en ajouter d’autres. Par exemple, comme lorsqu’on recherche “mountain” dans iPhoto et que des photos avec des montagnes apparaissent, ce serait intéressant de chercher dans plusieurs jeux de données avec un modèle multimodal comme CLIP
  • Il m’est arrivé de l’utiliser en expliquant une idée et un objectif, pour obtenir en retour le nom d’un domaine ou des mots-clés à rechercher
    Les grands modèles de langage ressemblent au meilleur moteur de recherche floue, et ils fonctionnent d’une manière assez singulière, mais complémentaire, par rapport aux moteurs de recherche traditionnels

• J’aime le concept, mais les résultats n’ont pas été très bons
  J’ai saisi “weird looking monkey” en m’attendant à des résultats comme le nasique ou le rhinopithèque de Roxellane, mais je n’ai obtenu que des articles comme “Pet monkey”, “List of individual monkeys”, “Ethnoprimatology”, “Monkey”
  En mettant la même requête dans Google, j’obtiens exactement les résultats attendus, donc c’était décevant ; je voulais découvrir des singes à l’apparence étrange que je ne connaissais pas

  • C’est un modèle sentence-transformer prêt à l’emploi datant de plus d’un an qui a été utilisé tel quel
    La démo visait surtout à montrer la base de données d’embeddings, mais les embeddings eux-mêmes sont aussi un peu utiles
    Comme je ne conserve absolument aucune donnée d’analyse sur ce que les gens trouvent ou ne trouvent pas sur la page, je ne suis pas vraiment prêt à améliorer les résultats de recherche
  • Les éditeurs de Wikipedia et ses recommandations n’aiment généralement pas les adjectifs subjectifs, et le “weird looking” de la requête ressemble à une formulation que les articles Wikipedia éviteraient
    Il devient donc difficile d’obtenir de bons résultats dans ce corpus de connaissances

• L’implémentation est vraiment bonne, et le fait de pouvoir faire ça hors ligne est impressionnant. En revanche, la qualité des embeddings semble encore insuffisante
  Une astuce qui pourrait aider serait de ne pas embarquer l’article Wikipedia entier, mais plutôt des phrases proches d’une définition, ou généralement la première phrase/le premier paragraphe. Je ne sais pas très bien quelle partie est utilisée actuellement
  Mon site OneLook propose aussi depuis 2003 une fonctionnalité similaire pour trouver des mots et des concepts à partir d’une description sur https://onelook.com/thesaurus/
  Au départ, c’était une pure recherche de dictionnaire inversé, mais au cours des 20 dernières années j’ai essayé les embeddings de mots, les embeddings de phrases, et plus récemment les grands modèles de langage ; aujourd’hui, GPT génère des candidats pour les entrées auxquelles le système ne peut pas répondre lui-même
  Pour cette tâche, les grands modèles de langage sont nettement meilleurs que les anciennes approches, ce qui a un peu réduit ma motivation à améliorer cette partie de OneLook. Je vois souvent des gens dire que la recherche par définition inversée est la principale raison pour laquelle ils utilisent ChatGPT

• Je vois ça un peu tard, mais les embeddings de texte, du moins ceux utilisés dans cet article, ne sont généralement pas très bons pour chercher par ambiance
  Ils consistent plutôt à comparer des mots qui se recoupent ou à trouver du contenu similaire à la requête
  Il existe toutefois un article récent qui s’est vraiment attaqué à ce problème : “Retrieving Texts based on Abstract Descriptions” (Ravfogel et al., 2023) https://arxiv.org/abs/2305.12517
  L’article contient de nombreux exemples de recherche à partir de descriptions abstraites, comme “un architecte qui conçoit des bâtiments”, “une entreprise qui fait partie d’une autre entreprise”, “un livre qui a influencé le développement d’un genre”
  Les embeddings en question semblent bien mieux prendre en charge ce type de recherche ; il serait donc intéressant de refaire la recherche Wikipedia hors ligne de l’article lié avec ce nouveau type d’embeddings

• Pour l’instant, la page ne fonctionne pas chez moi : model_quantized.onnx ne se charge pas
  Pendant que je tape, il a tout de même téléchargé jusqu’à 19,2 Mo à environ 50 Ko/s ; si chaque visiteur déclenche ça, cela pourrait faire des choses horribles aux frais de bande passante de Lee Butterman

  • C’est servi comme fichiers statiques depuis un seul t2.nano, donc je ne sais pas ce que ça va donner

• Ce que ça fait est très impressionnant, mais la qualité des résultats de recherche ne semble pas bonne
  Par expérience, je sais qu’il est vraiment difficile d’évaluer manuellement la qualité des résultats de recherche. On peut être très proche d’un excellent résultat tout en renvoyant une correspondance bien pire

  • Exact. La qualité ne sera probablement pas aussi bonne que Similar Website Finder https://explore2.marginalia.nu/ ;)
    Utiliser des embeddings de phrases plus récents devrait améliorer les résultats, et il faut que je collecte davantage de données

• La technologie est très impressionnante, mais les résultats ne l’étaient pas
  En cherchant « pointy building in Paris », j’ai obtenu 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
  Le bâtiment pointu le plus célèbre de Paris n’apparaissait pas du tout
  Pour ce genre d’application, les embeddings de phrases sur l’ensemble du document ne sont peut-être pas la meilleure option

  • Au moins 5 d’entre eux contiennent probablement la réponse recherchée
  • Si tu parles de la tour Eiffel, ce n’est pas un bâtiment
    Je viens de vérifier l’article : le mot « building » y apparaît 19 fois, mais la plupart du temps comme verbe, puis ensuite dans « Chrysler Building »
    Sauf s’il existe un autre bâtiment pointu célèbre auquel je ne pense pas

• Une partie de la magie des moteurs de recherche consiste à mélanger les embeddings des pages qui pointent vers cette page, ou des mots-clés de recherche d’information traditionnels, pondérés par les clics et les scores d’autorité
  Sans ce signal, beaucoup d’informations utiles sont ignorées, et les résultats semblent moins magiques
  Cela reste une démo impressionnante et intéressante

• J’aurais voulu aimer, mais dans les recherches que j’ai essayées, j’ai obtenu très peu de résultats pertinents
  Pour « The wizard in The Lord of the Rings », il n’y avait ni Gandalf ni Saruman, seulement des livres liés au Seigneur des anneaux
  Pour « Protagonist of Scorsese's Taxi Driver », il n’y avait pas Travis Bickle
  Pour « A person that plants trees for a living », bizarrement, gardener n’était pas dans la liste
  Pour « Curly-haired painter on TV », il n’y avait pas du tout Bob Ross
  Pour « Unusually shaped modern art museum in Spain », Bilbao apparaissait bien en 4e position, mais les autres n’avaient pas de forme particulière
  Pour « Dog shaped like a sausage », on s’attendrait à voir dachshund parmi les premiers résultats

  • Il faut noter que tous les résultats attendus ici ont une page Wikipedia
    S’ils n’avaient pas d’article, leur absence serait moins étrange, mais en réalité ils existent tous

• « Vibes » parle beaucoup plus que « sentence embeddings ». Je devrais peut-être commencer à utiliser cette expression moi aussi :)

  • C’est une expression qui sacrifie beaucoup de précision au profit de ce côté « parlant »
    L’auteur du billet n’a pas expliqué pourquoi il avait choisi ce mot, et cela ne correspond à aucun usage de « vibe » que je connaisse
    Je me demande si « gist » n’était pas assez à la mode