Comment ColBERT aide les développeurs à surmonter les limites de RAG

La génération augmentée par récupération (RAG) fait désormais partie intégrante des applications d’intelligence artificielle (IA) générative. Compléter votre invite de candidature avec un contexte pertinent extrait d'une base de données vectorielles peut considérablement augmenter la précision et réduire les hallucinations. Cela signifie que la pertinence croissante des résultats de recherche vectorielle est directement liée à la qualité de votre application RAG.

Il y a deux raisons pour lesquelles RAG reste populaire et de plus en plus pertinent même si les grands modèles de langage (LLM) augmentent leur fenêtre contextuelle :

1. Le temps de réponse et le prix du LLM augmentent tous deux linéairement avec la longueur du contexte.

2. Les LLM ont encore du mal à la fois à la récupération et au raisonnement dans des contextes massifs.

   
   

Mais RAG n’est pas une baguette magique. En particulier, la conception la plus courante, la récupération de passages denses (DPR), représente à la fois les requêtes et les passages comme un seul vecteur d'intégration et utilise une simple similarité cosinusoïdale pour évaluer la pertinence. Cela signifie que DPR s'appuie fortement sur le modèle d'intégration ayant l'étendue de la formation pour reconnaître tous les termes de recherche pertinents.

Malheureusement, les modèles disponibles dans le commerce sont confrontés à des termes inhabituels, notamment des noms, qui ne figurent pas couramment dans leurs données d'entraînement. DPR a également tendance à être hypersensible à la stratégie de segmentation, ce qui peut faire manquer un passage pertinent s'il est entouré de nombreuses informations non pertinentes. Tout cela impose au développeur d'applications de « réussir du premier coup », car une erreur entraîne généralement la nécessité de reconstruire l'index à partir de zéro.

Résoudre les défis de DPR avec ColBERT

ColBERT est une nouvelle façon d'évaluer la pertinence des passages en utilisant un modèle de langage BERT qui résout considérablement les problèmes liés au DPR. Ce diagramme du premier article de ColBERT montre pourquoi c'est si excitant :

Ceci compare les performances de ColBERT avec d'autres solutions de pointe pour l'ensemble de données MS-MARCO. (MS-MARCO est un ensemble de requêtes Bing pour lesquelles Microsoft a noté manuellement les passages les plus pertinents. C'est l'un des meilleurs critères de récupération.) En bas et à droite, c'est mieux.

En bref, ColBERT surpasse largement les solutions généralement beaucoup plus complexes au prix d'une légère augmentation de la latence.

Pour tester cela, j'ai créé une démo et indexé plus de 1 000 articles Wikipédia avec ada002 DPR et ColBERT. J'ai constaté que ColBERT fournit des résultats nettement meilleurs sur des termes de recherche inhabituels.

La capture d'écran suivante montre que DPR ne parvient pas à reconnaître le nom inhabituel de William H. Herndon, un associé d'Abraham Lincoln, tandis que ColBERT trouve la référence dans l'article de Springfield. Notez également que le résultat n°2 de ColBERT concerne un William différent, alors qu'aucun des résultats de DPR n'est pertinent.

ColBERT est souvent décrit dans un jargon dense d’apprentissage automatique, mais il est en réalité très simple. Je vais montrer comment implémenter la récupération et la notation ColBERT sur DataStax Astra DB avec seulement quelques lignes de Python et Cassandra Query Language (CQL).

La grande idée

Au lieu du DPR traditionnel basé sur un seul vecteur qui transforme les passages en un seul vecteur « d'intégration », ColBERT génère un vecteur influencé par le contexte pour chaque jeton dans les passages. ColBERT génère de la même manière des vecteurs pour chaque jeton de la requête.

(La tokenisation fait référence à la division de l'entrée en fractions de mots avant son traitement par un LLM. Andrej Karpathy, membre fondateur de l'équipe OpenAI, vient de publier une vidéo exceptionnelle sur la façon dont cela fonctionne .)

Ensuite, le score de chaque document est la somme de la similarité maximale de chaque intégration de requête avec l'une des intégrations de documents :

 def maxsim(qv, document_embeddings): return max(qv @ dv for dv in document_embeddings) def score(query_embeddings, document_embeddings): return sum(maxsim(qv, document_embeddings) for qv in query_embeddings)

(@ est l'opérateur PyTorch pour le produit scalaire et constitue la mesure la plus courante de similarité vectorielle .)

Voilà, vous pouvez implémenter la notation ColBERT dans quatre lignes de Python ! Vous comprenez désormais ColBERT mieux que 99 % des personnes qui publient des articles à ce sujet sur X (anciennement Twitter).

Le reste des articles de ColBERT traitent de :

1. Comment affiner le modèle BERT pour générer les meilleures intégrations pour un ensemble de données donné ?
2. Comment limiter l'ensemble de documents pour lesquels vous calculez le score (relativement coûteux) affiché ici ?

La première question est facultative et hors du cadre de cet article. J'utiliserai le point de contrôle ColBERT pré-entraîné. Mais la seconde est simple à réaliser avec une base de données vectorielle comme DataStax Astra DB.

ColBERT sur Astra DB

Il existe une bibliothèque Python tout-en-un populaire pour ColBERT appelée RAGatouille ; cependant, cela suppose un ensemble de données statique. L'une des fonctionnalités puissantes des applications RAG est de répondre aux données changeantes de manière dynamique en temps réel . À la place, je vais utiliser l'index vectoriel d'Astra pour restreindre l'ensemble des documents dont j'ai besoin pour classer les meilleurs candidats pour chaque sous-vecteur.

Il y a deux étapes lors de l'ajout de ColBERT à une application RAG : l'ingestion et la récupération.

Ingestion

Étant donné que plusieurs intégrations seront associées à chaque morceau de document, j'aurai besoin de deux tableaux :

 CREATE TABLE chunks ( title text, part int, body text, PRIMARY KEY (title, part) ); CREATE TABLE colbert_embeddings ( title text, part int, embedding_id int, bert_embedding vector<float, 128>, PRIMARY KEY (title, part, embedding_id) ); CREATE INDEX colbert_ann ON colbert_embeddings(bert_embedding) WITH OPTIONS = { 'similarity_function': 'DOT_PRODUCT' };

Après avoir installé la bibliothèque ColBERT ( pip install colbert-ai ) et téléchargé le point de contrôle BERT pré-entraîné , je peux charger des documents dans ces tables :

 from colbert.infra.config import ColBERTConfig from colbert.modeling.checkpoint import Checkpoint from colbert.indexing.collection_encoder import CollectionEncoder from cassandra.concurrent import execute_concurrent_with_args from db import DB def encode_and_save(title, passages): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encoder = CollectionEncoder(cf, cp) # encode_passages returns a flat list of embeddings and a list of how many correspond to each passage embeddings_flat, counts = encoder.encode_passages(passages) # split up embeddings_flat into a nested list start_indices = [0] + list(itertools.accumulate(counts[:-1])) embeddings_by_part = [embeddings_flat[start:start+count] for start, count in zip(start_indices, counts)] # insert into the database for part, embeddings in enumerate(embeddings_by_part): execute_concurrent_with_args(db.session, db.insert_colbert_stmt, [(title, part, i, e) for i, e in enumerate(embeddings)])

(J'aime encapsuler ma logique de base de données dans un module dédié ; vous pouvez accéder à la source complète dans mon référentiel GitHub .)

Récupération

La récupération ressemble alors à ceci :

 def retrieve_colbert(query): db = DB() cf = ColBERTConfig(checkpoint='checkpoints/colbertv2.0') cp = Checkpoint(cf.checkpoint, colbert_config=cf) encode = lambda q: cp.queryFromText([q])[0] query_encodings = encode(query) # find the most relevant documents for each query embedding. using a set # handles duplicates so we don't retrieve the same one more than once docparts = set() for qv in query_encodings: rows = db.session.execute(db.query_colbert_ann_stmt, [list(qv)]) docparts.update((row.title, row.part) for row in rows) # retrieve these relevant documents and score each one scores = {} for title, part in docparts: rows = db.session.execute(db.query_colbert_parts_stmt, [title, part]) embeddings_for_part = [tensor(row.bert_embedding) for row in rows] scores[(title, part)] = score(query_encodings, embeddings_for_part) # return the source chunk for the top 5 return sorted(scores, key=scores.get, reverse=True)[:5]

Voici la requête en cours d'exécution pour la partie des documents les plus pertinents ( db.query_colbert_ann_stmt ) :

 SELECT title, part FROM colbert_embeddings ORDER BY bert_embedding ANN OF ? LIMIT 5

Au-delà des bases : RAGStack

Cet article et le référentiel lié présentent brièvement le fonctionnement de ColBERT. Vous pouvez mettre en œuvre cela dès aujourd’hui avec vos propres données et voir des résultats immédiats. Comme pour tout ce qui concerne l’IA, les meilleures pratiques évoluent quotidiennement et de nouvelles techniques émergent constamment.

Pour faciliter le suivi de l'état de l'art, DataStax intègre cette amélioration ainsi que d'autres améliorations dans RAGStack , notre bibliothèque RAG prête pour la production exploitant LangChain et LlamaIndex. Notre objectif est de fournir aux développeurs une bibliothèque cohérente pour les applications RAG qui leur permet de contrôler la transition vers de nouvelles fonctionnalités. Au lieu d'avoir à suivre la myriade de changements dans les techniques et les bibliothèques, vous disposez d'un seul flux, vous pouvez donc vous concentrer sur la création de votre application. Vous pouvez utiliser RAGStack dès aujourd'hui pour intégrer les meilleures pratiques pour LangChain et LlamaIndex dès le départ ; des avancées comme ColBERT arriveront sur RAGstack dans les prochaines versions.

Par Jonathan Ellis, DataStax