ketannnn commited on
Commit
f42877a
·
1 Parent(s): 8c91ba1

fix(ml): restrict cross-encoder batch_size to prevent OOM kills on uncapped arrays

Browse files
Files changed (1) hide show
  1. backend/src/ml/reranker.py +2 -1
backend/src/ml/reranker.py CHANGED
@@ -17,7 +17,8 @@ def rerank(query: str, passages: list[str]) -> list[float]:
17
  return []
18
  reranker = _get_reranker()
19
  pairs = [[query, p] for p in passages]
20
- scores = reranker.compute_score(pairs, normalize=True)
 
21
  if isinstance(scores, float):
22
  scores = [scores]
23
  return [float(s) for s in scores]
 
17
  return []
18
  reranker = _get_reranker()
19
  pairs = [[query, p] for p in passages]
20
+ # Enforce small batch_size to prevent OOM kills on 100+ candidates in cloud environments
21
+ scores = reranker.compute_score(pairs, normalize=True, batch_size=8)
22
  if isinstance(scores, float):
23
  scores = [scores]
24
  return [float(s) for s in scores]