Recherche neuronale

Neural-Cherche est une bibliothèque conçue pour affiner les modèles de recherche neuronale tels que Splade, Colbert et Sparseembed sur un ensemble de données spécifique. Neural-Cherche fournit également des classes pour exécuter une inférence efficace sur un retriever ou un rang affiné. Neural-Cherche vise à offrir une méthode simple et efficace pour le réglage fin et l'utilisation de modèles de recherche neuronale dans les paramètres hors ligne et en ligne. Il permet également aux utilisateurs d'enregistrer toutes les intégres calculés pour éviter les calculs redondants.
Le neural-cherche est compatible avec les dispositifs CPU, GPU et MPS. Nous pouvons affiner Colbert à partir de tout point de contrôle pré-formé Transformateur de phrase. Le déplatement et le SPARseembed sont plus délicats à affiner et ont besoin d'un modèle pré-formé MLM.
Nous pouvons installer neural-cherche en utilisant:
pip install neural-cherche
Si nous prévoyons d'évaluer notre modèle pendant l'installation de la formation:
pip install "neural-cherche[eval]"
La documentation complète est disponible ici.
Votre ensemble de données de formation doit être fait de triplets (anchor, positive, negative) où l'ancre est une requête, positive est un document directement lié à l'ancre et négatif est un document qui n'est pas pertinent pour l'ancre.
X = [
( "anchor 1" , "positive 1" , "negative 1" ),
( "anchor 2" , "positive 2" , "negative 2" ),
( "anchor 3" , "positive 3" , "negative 3" ),
]Et voici comment affiner Colbert à partir d'un point de contrôle pré-formé Transformateur de phrase à l'aide de neural-cherche:
import torch
from neural_cherche import models , utils , train
model = models . ColBERT (
model_name_or_path = "raphaelsty/neural-cherche-colbert" ,
device = "cuda" if torch . cuda . is_available () else "cpu" # or mps
)
optimizer = torch . optim . AdamW ( model . parameters (), lr = 3e-6 )
X = [
( "query" , "positive document" , "negative document" ),
( "query" , "positive document" , "negative document" ),
( "query" , "positive document" , "negative document" ),
]
for step , ( anchor , positive , negative ) in enumerate ( utils . iter (
X ,
epochs = 1 , # number of epochs
batch_size = 8 , # number of triples per batch
shuffle = True
)):
loss = train . train_colbert (
model = model ,
optimizer = optimizer ,
anchor = anchor ,
positive = positive ,
negative = negative ,
step = step ,
gradient_accumulation_steps = 50 ,
)
if ( step + 1 ) % 1000 == 0 :
# Save the model every 1000 steps
model . save_pretrained ( "checkpoint" )Voici comment utiliser le modèle Colbert affiné pour re-rier les documents:
import torch
from lenlp import sparse
from neural_cherche import models , rank , retrieve
documents = [
{ "id" : "doc1" , "title" : "Paris" , "text" : "Paris is the capital of France." },
{ "id" : "doc2" , "title" : "Montreal" , "text" : "Montreal is the largest city in Quebec." },
{ "id" : "doc3" , "title" : "Bordeaux" , "text" : "Bordeaux in Southwestern France." },
]
retriever = retrieve . BM25 (
key = "id" ,
on = [ "title" , "text" ],
count_vectorizer = sparse . CountVectorizer (
normalize = True , ngram_range = ( 3 , 5 ), analyzer = "char_wb" , stop_words = []
),
k1 = 1.5 ,
b = 0.75 ,
epsilon = 0.0 ,
)
model = models . ColBERT (
model_name_or_path = "raphaelsty/neural-cherche-colbert" ,
device = "cuda" if torch . cuda . is_available () else "cpu" , # or mps
)
ranker = rank . ColBERT (
key = "id" ,
on = [ "title" , "text" ],
model = model ,
)
documents_embeddings = retriever . encode_documents (
documents = documents ,
)
retriever . add (
documents_embeddings = documents_embeddings ,
)Nous pouvons maintenant récupérer des documents en utilisant le modèle affiné:
queries = [ "Paris" , "Montreal" , "Bordeaux" ]
queries_embeddings = retriever . encode_queries (
queries = queries ,
)
ranker_queries_embeddings = ranker . encode_queries (
queries = queries ,
)
candidates = retriever (
queries_embeddings = queries_embeddings ,
batch_size = 32 ,
k = 100 , # number of documents to retrieve
)
# Compute embeddings of the candidates with the ranker model.
# Note, we could also pre-compute all the embeddings.
ranker_documents_embeddings = ranker . encode_candidates_documents (
candidates = candidates ,
documents = documents ,
batch_size = 32 ,
)
scores = ranker (
queries_embeddings = ranker_queries_embeddings ,
documents_embeddings = ranker_documents_embeddings ,
documents = candidates ,
batch_size = 32 ,
)
scores [[{ 'id' : 0 , 'similarity' : 22.825355529785156 },
{ 'id' : 1 , 'similarity' : 11.201947212219238 },
{ 'id' : 2 , 'similarity' : 10.748161315917969 }],
[{ 'id' : 1 , 'similarity' : 23.21628189086914 },
{ 'id' : 0 , 'similarity' : 9.9658203125 },
{ 'id' : 2 , 'similarity' : 7.308732509613037 }],
[{ 'id' : 1 , 'similarity' : 6.4031805992126465 },
{ 'id' : 0 , 'similarity' : 5.601611137390137 },
{ 'id' : 2 , 'similarity' : 5.599479675292969 }]] Neural-Cherche fournit un SparseEmbed , un SPLADE , un TFIDF , un Retriever BM25 et un ColBERT Ranker qui peut être utilisé pour réorganiser la sortie d'un retriever. Pour plus d'informations, veuillez vous référer à la documentation.
Nous fournissons des points de contrôle pré-formés spécifiquement conçus pour le neural-cherche: Raphaelsty / Neural-Cherche-Sparse-Embed et Raphaelsty / Neural-Cherche-Colbert. Ces points de contrôle sont affinés sur un sous-ensemble de l'ensemble de données MS-MARCO et bénéficieraient d'une affection affinée sur votre ensemble de données spécifique. Vous pouvez affiner Colbert à partir de n'importe quel point de contrôle pré-formé Transformateur de phrase afin d'adapter votre langue spécifique. Vous devez utiliser un point de vérification basé sur MLM pour affiner SPARseembed.
| ensemble de données scifact | ||||
|---|---|---|---|---|
| modèle | Point de contrôle HuggingFace | ndcg @ 10 | frappe @ 10 | frappe @ 1 |
| Tfidf | - | 0,62 | 0,86 | 0,50 |
| BM25 | - | 0,69 | 0,92 | 0,56 |
| Clairsemé | Raphaelsty / Neural-Cherche-Spares | 0,62 | 0,87 | 0,48 |
| Transformateur de phrase | Transformateurs de phrase / All-Mpnet-Base-V2 | 0,66 | 0,89 | 0,53 |
| Colbert | Raphaelsty / Neural-Cherche-Colbert | 0,70 | 0,92 | 0,58 |
| Tfidf Retriever + Colbert Ranker | Raphaelsty / Neural-Cherche-Colbert | 0,71 | 0,94 | 0,59 |
| BM25 Retriever + Colbert Ranker | Raphaelsty / Neural-Cherche-Colbert | 0,72 | 0,95 | 0,59 |
Salle: Modèle de lexical et d'expansion clairsemé pour le classement de première étape rédigé par Thibault Foral, Benjamin Piwowarski, Stéphane Clinchant, Sigir 2021.
Salle V2: Modèle de lexical et d'expansion clairsemé pour la recherche d'informations rédigée par Thibault Foral, Carlos Lassance, Benjamin Piwowarski, Stéphane Clinchant, Sigir 2022.
Sparesembed: Apprentissage des représentations lexicales clairsemées avec des incorporations contextuelles pour la récupération rédigées par Weize Kong, Jeffrey M. Dudek, Cheng Li, Mingyang Zhang et Mike Bendersky, Sigir 2023.
Colbert: Recherche de passage efficace et efficace via une interaction tardive contextualisée sur Bert rédigée par Omar Khattab, Matei Zaharia, Sigir 2020.
Cette bibliothèque Python est sous licence sous la licence open source du MIT, et le modèle de volet est sous licence non commerciale uniquement par les auteurs. Sparesembed et Colbert sont entièrement open-source, y compris l'utilisation commerciale.