Package: neural-tree
Neural-TreeNeural SearchAre tree-based indexes the counterpart of standard ANN algorithms for token-level embeddings IR models? Neural-Tree replicate the SIGIR 2023 publicationConstructing Tree-based Index for Efficient and Effective Dense Retrievalin order to accelerate ColBERT. Neural-Tree is compatible with Sentence Transformers and TfIdf models as in the original paper.Neural-Tree creates a tree using hierarchical clustering of documents and then learn embeddings in each node of the tree using paired queries and documents. Additionally, there is the flexibility to input an existing tree structure in JSON format to build the index.The optimization of the index by Neural-Tree is geared towards maintaining the performance level of the original model while significantly speeding up the search process. It is important to note that Neural-Tree does not modify the underlying model; therefore, it is advisable to initiate tree creation with a model that has already been fine-tuned. Given that Neural-Tree does not alter the model, the index training process is relatively quick.InstallationWe can install neural-tree using:pip install neural-treeIf we plan to evaluate our model while training install:pip install "neural-tree[eval]"DocumentationThe complete documentation is availablehere.Quick StartThe following code shows how to train a tree index. Let's start by creating a fictional dataset:documents=[{"id":0,"content":"paris"},{"id":1,"content":"london"},{"id":2,"content":"berlin"},{"id":3,"content":"rome"},{"id":4,"content":"bordeaux"},{"id":5,"content":"milan"},]train_queries=["paris is the capital of france","london is the capital of england","berlin is the capital of germany","rome is the capital of italy",]train_documents=[{"id":0,"content":"paris"},{"id":1,"content":"london"},{"id":2,"content":"berlin"},{"id":3,"content":"rome"},]test_queries=["bordeaux is the capital of france","milan is the capital of italy",]Let's train the index using thedocuments,train_queriesandtrain_documentswe have gathered.importtorchfromneural_chercheimportmodelsfromneural_treeimportclustering,trees,utilsmodel=models.ColBERT(model_name_or_path="raphaelsty/neural-cherche-colbert",device="cuda"iftorch.cuda.is_available()else"cpu",)tree=trees.ColBERT(key="id",on=["content"],model=model,documents=documents,leaf_balance_factor=100,# Number of documents per leafbranch_balance_factor=5,# Number children per noden_jobs=-1,# set to 1 with Google Colab)optimizer=torch.optim.AdamW(lr=3e-3,params=list(tree.parameters()))forstep,batch_queries,batch_documentsinutils.iter(queries=train_queries,documents=train_documents,shuffle=True,epochs=50,batch_size=32,):loss=tree.loss(queries=batch_queries,documents=batch_documents,)loss.backward()optimizer.step()optimizer.zero_grad(set_to_none=True)Let's now duplicate some documents of the tree in order to increase accuracy.documents_to_leafs=clustering.optimize_leafs(tree=tree,queries=train_queries+test_queries,documents=documents,)tree=tree.add(documents=documents,documents_to_leafs=documents_to_leafs,)We are now ready to retrieve documents:scores=tree(queries=["bordeaux","milan"],k_leafs=2,k=2,)print(scores["documents"])[[{"id":4,"similarity":5.28,"leaf":"12"},{"id":0,"similarity":3.17,"leaf":"12"},],[{"id":5,"similarity":5.11,"leaf":"10"},{"id":2,"similarity":3.57,"leaf":"10"},],]EvaluationWe can evaluate the performance of the tree using the following code:documents,queries_ids,test_queries,qrels=datasets.load_beir_test(dataset_name="scifact",)candidates=tree(queries=test_queries,k_leafs=2,k=10,)scores=utils.evaluate(scores=candidates["documents"],qrels=qrels,queries_ids=queries_ids,)print(scores)BenchmarksScifact DatasetVanillaNeural-TreemodelHuggingFace Checkpointndcg@10hits@10hits@1queries / secondndcg@10hits@10hits@1queries / secondAccelerationTfIdfCherche-0,610,850,477600,560,820,421080+42.11%SentenceTransformer GPUFaiss.IndexFlatL2 CPUsentence-transformers/all-mpnet-base-v20,660,890,534750,660,880,53518+9.05%ColBERTNeural-Cherche GPUraphaelsty/neural-cherche-colbert0,700,920,5830,700,910,59256x85Note that this benchmark do not implementColBERTV2efficient retrieval but rather compare ColBERT raw retrieval with Neural-Tree. We could accelerate SentenceTransformer vanilla by using optimized Faiss index.ContributingWe welcome contributions to Neural-Tree, a tool designed to enhance tree visualization, model node topics, and leverage the tree structure to expedite Large Language Model (LLM) searches. Our focus includes refining the clustering of ColBERT embeddings through hierarchical clustering, which is currently facilitated by TfIdf. Additionally, there's an opportunity to contribute towards optimizing clustering, aiming to achieve comprehensive ColBERT cluster optimization independently of TfIdf.LicenseThis project is licensed under the terms of the MIT license.ReferencesConstructing Tree-based Index for Efficient and Effective Dense Retrieval, GithubColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERTMyriade