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:
@@ -487,6 +487,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
|
||||
# Reranking settings
|
||||
'RAG_RERANKING_MODEL': request.app.state.config.RAG_RERANKING_MODEL,
|
||||
'RAG_RERANKING_ENGINE': request.app.state.config.RAG_RERANKING_ENGINE,
|
||||
'RAG_RERANKING_BATCH_SIZE': request.app.state.config.RAG_RERANKING_BATCH_SIZE,
|
||||
'RAG_EXTERNAL_RERANKER_URL': request.app.state.config.RAG_EXTERNAL_RERANKER_URL,
|
||||
'RAG_EXTERNAL_RERANKER_API_KEY': request.app.state.config.RAG_EXTERNAL_RERANKER_API_KEY,
|
||||
'RAG_EXTERNAL_RERANKER_TIMEOUT': request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT,
|
||||
@@ -694,6 +695,7 @@ class ConfigForm(BaseModel):
|
||||
# Reranking settings
|
||||
RAG_RERANKING_MODEL: Optional[str] = None
|
||||
RAG_RERANKING_ENGINE: Optional[str] = None
|
||||
RAG_RERANKING_BATCH_SIZE: Optional[int] = None
|
||||
RAG_EXTERNAL_RERANKER_URL: Optional[str] = None
|
||||
RAG_EXTERNAL_RERANKER_API_KEY: Optional[str] = None
|
||||
RAG_EXTERNAL_RERANKER_TIMEOUT: Optional[str] = None
|
||||
@@ -940,6 +942,12 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
||||
else request.app.state.config.RAG_EXTERNAL_RERANKER_TIMEOUT
|
||||
)
|
||||
|
||||
request.app.state.config.RAG_RERANKING_BATCH_SIZE = (
|
||||
form_data.RAG_RERANKING_BATCH_SIZE
|
||||
if form_data.RAG_RERANKING_BATCH_SIZE is not None
|
||||
else request.app.state.config.RAG_RERANKING_BATCH_SIZE
|
||||
)
|
||||
|
||||
log.info(
|
||||
f'Updating reranking model: {request.app.state.config.RAG_RERANKING_MODEL} to {form_data.RAG_RERANKING_MODEL}'
|
||||
)
|
||||
@@ -967,6 +975,7 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
|
||||
request.app.state.config.RAG_RERANKING_ENGINE,
|
||||
request.app.state.config.RAG_RERANKING_MODEL,
|
||||
request.app.state.rf,
|
||||
reranking_batch_size=request.app.state.config.RAG_RERANKING_BATCH_SIZE,
|
||||
)
|
||||
except Exception as e:
|
||||
log.error(f'Error loading reranking model: {e}')
|
||||
|
||||
Reference in New Issue
Block a user