Package: neural-cherche
Neural-ChercheNeural SearchNeural-Cherche is a library designed to fine-tune neural search models such as Splade, ColBERT, and SparseEmbed on a specific dataset. Neural-Cherche also provide classes to run efficient inference on a fine-tuned retriever or ranker. Neural-Cherche aims to offer a straightforward and effective method for fine-tuning and utilizing neural search models in both offline and online settings. It also enables users to save all computed embeddings to prevent redundant computations.Neural-Cherche is compatible with CPU, GPU and MPS devices. We can fine-tune ColBERT from any
Sentence Transformer pre-trained checkpoint. Splade and SparseEmbed are more tricky to fine-tune and need a MLM pre-trained model.InstallationWe can install neural-cherche using:pip install neural-chercheIf we plan to evaluate our model while training install:pip install "neural-cherche[eval]"DocumentationThe complete documentation is availablehere.Quick StartYour training dataset must be made out of triples(anchor, positive, negative)where anchor is a query, positive is a document that is directly linked to the anchor and negative is a document that is not relevant for the anchor.X=[("anchor 1","positive 1","negative 1"),("anchor 2","positive 2","negative 2"),("anchor 3","positive 3","negative 3"),]And here is how to fine-tune ColBERT from a Sentence Transformer pre-trained checkpoint using neural-cherche:importtorchfromneural_chercheimportmodels,utils,trainmodel=models.ColBERT(model_name_or_path="raphaelsty/neural-cherche-colbert",device="cuda"iftorch.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"),]forstep,(anchor,positive,negative)inenumerate(utils.iter(X,epochs=1,# number of epochsbatch_size=8,# number of triples per batchshuffle=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 stepsmodel.save_pretrained("checkpoint")RetrievalHere is how to use the fine-tuned ColBERT model to retrieve documents:fromneural_chercheimportmodels,retrieveimporttorchbatch_size=32documents=[{"id":0,"document":"Food"},{"id":1,"document":"Sports"},{"id":2,"document":"Cinema"},]model=models.ColBERT(model_name_or_path="raphaelsty/neural-cherche-colbert",device="cuda"iftorch.cuda.is_available()else"cpu",# or mps)retriever=retrieve.ColBERT(key="id",on=["document"],# the field to search on, can be a list of fieldsmodel=model,)documents_embeddings=retriever.encode_documents(documents=documents,batch_size=batch_size,)retriever=retriever.add(documents_embeddings=documents_embeddings,)Now we can retrieve documents using the fine-tuned model:queries_embeddings=retriever.encode_queries(queries=["Food","Sports","Cinema"],# list of queriesbatch_size=batch_size,)scores=retriever(queries_embeddings=queries_embeddings,batch_size=batch_size,k=10,# number of documents to retrieve)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}]]Please note that neural-cherche provide documentation to useColBERT as a rankerwhich is much more efficient.Neural-Cherche also provides aSparseEmbed, aSPLADE, aTFIDFretriever and aColBERTranker which can be used to re-order output of a retriever. For more information, please refer to thedocumentation.Pre-trained ModelsWe provide pre-trained checkpoints specifically designed for neural-cherche:raphaelsty/neural-cherche-sparse-embedandraphaelsty/neural-cherche-colbert. Those checkpoints are fine-tuned on a subset of the MS-MARCO dataset and would benefit from being fine-tuned on your specific dataset. You can fine-tune ColBERT from any Sentence Transformer pre-trained checkpoint in order to fit your specific language. You should use a MLM based-checkpoint to fine-tune SparseEmbed.scifact datasetmodelHuggingFace Checkpointndcg@10hits@10hits@1TfIdf-0,610,850,47SparseEmbedraphaelsty/neural-cherche-sparse-embed0,620,870,48Sentence Transformersentence-transformers/all-mpnet-base-v20,660,890,53ColBERTraphaelsty/neural-cherche-colbert0,700,920,58TfIDF Retriever + ColBERT Rankerraphaelsty/neural-cherche-colbert0,710,940,59Neural-Cherche ContributorsBenjamin ClaviéReferencesSPLADE: Sparse Lexical and Expansion Model for First Stage Rankingauthored by Thibault Formal, Benjamin Piwowarski, Stéphane Clinchant, SIGIR 2021.SPLADE v2: Sparse Lexical and Expansion Model for Information Retrievalauthored by Thibault Formal, Carlos Lassance, Benjamin Piwowarski, Stéphane Clinchant, SIGIR 2022.SparseEmbed: Learning Sparse Lexical Representations with Contextual Embeddings for Retrievalauthored by Weize Kong, Jeffrey M. Dudek, Cheng Li, Mingyang Zhang, and Mike Bendersky, SIGIR 2023.ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERTauthored by Omar Khattab, Matei Zaharia, SIGIR 2020.LicenseThis Python library is licensed under the MIT open-source license, and the splade model is licensed as non-commercial only by the authors. SparseEmbed and ColBERT are fully open-source including commercial usage.