feat: add RAG_RERANKING_BATCH_SIZE configuration option
Add configurable reranker batch size (env var RAG_RERANKING_BATCH_SIZE, default 32) following the same pattern as RAG_EMBEDDING_BATCH_SIZE. - config.py: PersistentConfig for RAG_RERANKING_BATCH_SIZE - main.py: import, state init, pass to get_reranking_function - colbert.py: accept batch_size param in predict() (was hardcoded 32) - utils.py: get_reranking_function passes batch_size at call time - retrieval.py: expose in config GET/POST endpoints and ConfigForm - Documents.svelte: add Reranking Batch Size input in admin settings Closes #23730
This commit is contained in:
@@ -59,14 +59,14 @@ class ColBERT(BaseReranker):
|
||||
|
||||
return normalized_scores.detach().cpu().numpy().astype(np.float32)
|
||||
|
||||
def predict(self, sentences):
|
||||
def predict(self, sentences, batch_size=32):
|
||||
query = sentences[0][0]
|
||||
docs = [i[1] for i in sentences]
|
||||
|
||||
# Embedding the documents
|
||||
embedded_docs = self.ckpt.docFromText(docs, bsize=32)[0]
|
||||
embedded_docs = self.ckpt.docFromText(docs, bsize=batch_size)[0]
|
||||
# Embedding the queries
|
||||
embedded_queries = self.ckpt.queryFromText([query], bsize=32)
|
||||
embedded_queries = self.ckpt.queryFromText([query], bsize=batch_size)
|
||||
embedded_query = embedded_queries[0]
|
||||
|
||||
# Calculate retrieval scores for the query against all documents
|
||||
|
||||
Reference in New Issue
Block a user