performance: Optimize ColBERT index free search with torch.topk

[Diegi97]

The following line is a bottleneck when using the index free search:

https://github.com/bclavie/RAGatouille/blob/796b49388cad0822564f30ef0b1464d021186637/ragatouille/models/colbert.py#L480

The change I introduce reduces the search time over 25k documents from 5.571s to 0.023s in my local setup with an intel i7 and a RTX 4090.

Script to reproduce:

from ragatouille import RAGPretrainedModel
from datasets import load_dataset

# Load the pretrained model
r = RAGPretrainedModel.from_pretrained('colbert-ir/colbertv2.0')

# Load the dataset
dataset = load_dataset('mteb/scidocs', 'corpus')
docs = dataset['corpus']['text']
print(f"Number of documents: {len(docs)}")

# Encode the documents
encodings = r.encode(docs, bsize=256)

# Perform search on encoded documents
import timeit

def search():
    return r.search_encoded_docs('Recurrent Neural Networks', k=5)

# Timing the searches
rnn_time = timeit.timeit(search, number=7)

print(f"Search 'Recurrent Neural Networks': {rnn_time / 7:.3f} s per loop")

[bclavie]

Thanks a lot! Let this be the first merge of the RAGatouille long overdue overhaul 😄