performance: Optimize ColBERT index free search with torch.topk

[Diegi97]

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

RAGatouille/ragatouille/models/colbert.py

Line 480 in 796b493

sorted_scores = sorted(enumerate(scores), key=lambda x: x[1], reverse=True)

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")