Búsqueda neuronal

Neural-Cherche es una biblioteca diseñada para ajustar modelos de búsqueda neuronal como Splade, Colbert y Sparseembed en un conjunto de datos específico. Neural-Cherche también proporciona clases para ejecutar una inferencia eficiente en un retriever o ranker sintonizado. Neural-Cherche tiene como objetivo ofrecer un método sencillo y efectivo para ajustar y utilizar modelos de búsqueda neuronal tanto en entornos fuera de línea como en línea. También permite a los usuarios guardar todos los incrustaciones calculadas para evitar cálculos redundantes.
Neural-Cherche es compatible con los dispositivos CPU, GPU y MPS. Podemos ajustar a Colbert desde cualquier punto de control previamente capacitado del transformador de oraciones. Flade y Sparseembed son más difíciles de ajustar y necesitan un modelo pretrontrado MLM.
Podemos instalar la clandestina neural usando:
pip install neural-cherche
Si planeamos evaluar nuestro modelo mientras la instalación de capacitación:
pip install "neural-cherche[eval]"
La documentación completa está disponible aquí.
Su conjunto de datos de entrenamiento debe estar hecho de triples (anchor, positive, negative) donde el ancla es una consulta, positivo es un documento que está directamente vinculado al ancla y negativo es un documento que no es relevante para el ancla.
X = [
( "anchor 1" , "positive 1" , "negative 1" ),
( "anchor 2" , "positive 2" , "negative 2" ),
( "anchor 3" , "positive 3" , "negative 3" ),
]Y aquí está cómo ajustar a Colbert a partir de un punto de control previamente capacitado del transformador de oraciones utilizando la clerín neuronal:
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" )Aquí le mostramos cómo usar el modelo Colbert ajustado para volver a clasificar los documentos:
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 ,
)Ahora podemos recuperar documentos utilizando el modelo ajustado:
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 proporciona un SparseEmbed , un SPLADE , un TFIDF , un BM25 Retriever y un Ranker ColBERT que puede usarse para reordenar la salida de un Retriever. Para obtener más información, consulte la documentación.
Proporcionamos puntos de control previamente capacitados diseñados específicamente para la clerín neuronal: Raphaelsty/Neural-Cherche-Sparse-Embed y Raphaelsty/Neural-Cherche-Colbert. Esos puntos de control están ajustados en un subconjunto del conjunto de datos MS-Marco y se beneficiarían de ser ajustados en su conjunto de datos específico. Puede ajustar a Colbert desde cualquier punto de control previamente capacitado del transformador de oraciones para que se ajuste a su idioma específico. Debe usar un punto de verificación basado en MLM para ajustar el fino.
| conjunto de datos de sicactación | ||||
|---|---|---|---|---|
| modelo | Punto de control de la cara de abrazos | ndcg@10 | golpes@10 | golpes@1 |
| Tfidf | - | 0,62 | 0,86 | 0,50 |
| BM25 | - | 0,69 | 0,92 | 0,56 |
| Escasez | Raphaelsty/Neural-Cherche-Sparse-Embed | 0,62 | 0,87 | 0,48 |
| Transformador de oraciones | oración-transformadores/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 |
Flade: modelo de expansión y léxico escaso para la clasificación de la primera etapa escrita por Thibault Formal, Benjamin Piwowarski, Stéphane Clinchant, Sigir 2021.
Splade V2: modelo de expansión y léxico escaso para la recuperación de información escrita por Thibault Formal, Carlos Lassance, Benjamin Piwowarski, Stéphane Clinchant, Sigir 2022.
SPARSEEMBED: Aprender representaciones léxicas escasas con integridades contextuales para la recuperación escrita por Weize Kong, Jeffrey M. Dudek, Cheng Li, Mingyang Zhang y Mike Bendersky, Sigir 2023.
Colbert: búsqueda eficiente y efectiva de pasaje a través de la interacción tardía contextualizada sobre Bert, escrito por Omar Khattab, Matei Zaharia, Sigir 2020.
Esta biblioteca de Python tiene licencia bajo la licencia de código abierto del MIT, y el modelo Splade tiene licencia como no comercial solo por los autores. Sparseembed y Colbert son de código abierto, incluido el uso comercial.