From 622dcc68e9f73cf579cac0d216f17f731b38d780 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:04:17 -0500 Subject: [PATCH 1/4] fix: Bound reranker input length and attention memory per batch --- tests/test-model-microbatches-735.py | 2 + tests/test-reranker-token-budget.py | 129 +++++++++++++++++++++++++++ truememory/model_server.py | 24 ++++- truememory/rerank_budget.py | 85 ++++++++++++++++++ truememory/reranker.py | 4 +- 5 files changed, 239 insertions(+), 5 deletions(-) create mode 100644 tests/test-reranker-token-budget.py create mode 100644 truememory/rerank_budget.py diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py index a241ab11..0d92d236 100644 --- a/tests/test-model-microbatches-735.py +++ b/tests/test-model-microbatches-735.py @@ -628,6 +628,8 @@ def safe_import(name: str, *args: object, **kwargs: object) -> object: return mps if name == "truememory.reranker": return types.SimpleNamespace(get_current_reranker_name=lambda: "Alibaba-NLP/gte-reranker-modernbert-base") + if name == "truememory.rerank_budget": + return builtins.__import__(name, *args, **kwargs) raise AssertionError("Unexpected runtime import: " + name) numpy, clock = OrderNP(), FakeClock() diff --git a/tests/test-reranker-token-budget.py b/tests/test-reranker-token-budget.py new file mode 100644 index 00000000..62845941 --- /dev/null +++ b/tests/test-reranker-token-budget.py @@ -0,0 +1,129 @@ +"""Reranker input cap and token-budgeted batches bound attention scratch.""" +from __future__ import annotations + +import numpy as np +import pytest + +from truememory import rerank_budget +from truememory.rerank_budget import ( + RERANK_MAX_TOKENS, RERANK_TOKEN_BUDGET, budget_batch_size, budgeted_predict, + budgeted_slices, cap_reranker_length, pair_token_lengths, rerank_token_budget, +) + + +class FakeCrossEncoder: + """Word count stands in for tokens; scores depend only on each pair.""" + + def __init__(self, max_seq_length: int | None = 1024) -> None: + self.max_seq_length = max_seq_length + self.calls: list[tuple[int, int]] = [] + + def tokenizer(self, first, second, truncation, max_length): + assert truncation == "longest_first" + return {"input_ids": [[0] * min(len(a.split()) + len(b.split()) + 3, max_length) + for a, b in zip(first, second)]} + + @staticmethod + def _input_length(pair) -> int: + return len(pair[0]) + len(pair[1]) + + def predict(self, pairs, batch_size, show_progress_bar=False): + self.calls.append((len(pairs), batch_size)) + return np.array([len(b.split()) / 1000 + len(a) / 1e6 for a, b in pairs], dtype=np.float32) + + +def _pairs(*word_counts: int) -> list[list[str]]: + return [["synthetic query", " ".join(["word"] * count)] for count in word_counts] + + +def test_cap_lowers_long_limits_only(): + for current, expected in ((8192, RERANK_MAX_TOKENS), (None, RERANK_MAX_TOKENS), (512, 512)): + model = FakeCrossEncoder(current) + cap_reranker_length(model) + assert model.max_seq_length == expected + + +def test_budget_shrinks_under_memory_pressure(): + assert rerank_token_budget(16 * 1024 ** 3) == RERANK_TOKEN_BUDGET + assert rerank_token_budget(4 * 1024 ** 3) == RERANK_TOKEN_BUDGET // 2 + assert rerank_token_budget(1024 ** 3) == RERANK_TOKEN_BUDGET // 4 + + +def test_default_batches_kept_for_pairs_up_to_512_tokens(): + assert budget_batch_size(512, 64, RERANK_TOKEN_BUDGET) == 64 + assert budgeted_slices([512, 40, 300] * 30, 64, RERANK_TOKEN_BUDGET) is None + + +def test_long_pairs_go_longest_first_in_budgeted_slices(): + lengths = [50] * 80 + [1024] * 20 + slices = budgeted_slices(lengths, 64, RERANK_TOKEN_BUDGET) + assert sorted(index for group in slices for index in group) == list(range(100)) + for group in slices: + assert len(group) * max(lengths[index] for index in group) ** 2 <= RERANK_TOKEN_BUDGET or len(group) == 1 + assert len(group) <= 64 + assert [len(group) for group in slices] == [16, 16, 64, 4] + + +def test_short_search_is_one_unchanged_predict_call(): + model = FakeCrossEncoder() + pairs = _pairs(*range(10, 110)) + scores = budgeted_predict(model, pairs, 64) + assert model.calls == [(100, 64)] + assert scores.shape == (100,) + + +def test_long_search_scores_match_and_respect_budget(monkeypatch): + monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET) + model = FakeCrossEncoder() + pairs = _pairs(*([20] * 70 + [2000] * 30)) + scores = budgeted_predict(model, pairs, 64) + expected = FakeCrossEncoder().predict(pairs, batch_size=len(pairs)) + assert np.array_equal(scores, expected) + assert sum(size for size, _ in model.calls) == 100 + lengths = pair_token_lengths(model, pairs) + assert max(lengths) == RERANK_MAX_TOKENS + assert all(batch == size and size <= 64 for size, batch in model.calls) + assert model.calls[0][0] == RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2 + + +def test_unmeasurable_model_keeps_single_call(): + class Proxy: + calls = [] + + def predict(self, pairs, batch_size, show_progress_bar=False): + self.calls.append(len(pairs)) + return np.zeros(len(pairs), dtype=np.float32) + + proxy = Proxy() + budgeted_predict(proxy, _pairs(5000, 5000), 64) + assert proxy.calls == [2] + + +def test_model_server_rerank_uses_budgeted_slices(monkeypatch): + from truememory.model_server import ModelServer + + fake = FakeCrossEncoder() + monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake) + monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET) + pairs = _pairs(*([30] * 50 + [1500] * 10)) + response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"}) + assert response["ok"] is True + expected = FakeCrossEncoder().predict(pairs, batch_size=len(pairs)) + assert np.allclose(response["scores"], expected, rtol=0, atol=0) + assert max(size for size, _ in fake.calls) <= 64 + assert fake.calls[0] == (RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2, RERANK_TOKEN_BUDGET // RERANK_MAX_TOKENS ** 2) + + +def test_model_server_short_rerank_unchanged(monkeypatch): + from truememory.model_server import ModelServer + + fake = FakeCrossEncoder() + monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake) + response = ModelServer().handle_request({"op": "rerank", "pairs": _pairs(*range(5, 45)), "model_name": "m"}) + assert response["ok"] is True + assert fake.calls == [(40, 64)] + + +@pytest.mark.parametrize("length", [1, 512, 513, 1024]) +def test_batch_size_never_zero_or_above_limit(length): + assert 1 <= budget_batch_size(length, 64, RERANK_TOKEN_BUDGET) <= 64 diff --git a/truememory/model_server.py b/truememory/model_server.py index 56f018bf..8ec53bf3 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -1078,6 +1078,8 @@ def _get_reranker( self._memory_stage("reranker construction", deadline, model_identity=(name, 0, "reranker"), device=device) model = CrossEncoder(name, device=device) + from truememory.rerank_budget import cap_reranker_length + cap_reranker_length(model) self._check_process_memory("reranker construction completion", deadline) self._reranker = model self._reranker_name = name @@ -1428,6 +1430,7 @@ def _handle_request_inner(self, request: dict) -> dict: scores = None offset = 0 order = None + steps = None recovery_name = model_name while offset < len(pairs) or scores is None: retry = None @@ -1441,18 +1444,31 @@ def _handle_request_inner(self, request: dict) -> dict: order = self._rerank_global_order(reranker, pairs, limit, deadline) if order is not None: recovery_name = self._reranker_name + from truememory.rerank_budget import budgeted_slices, pair_token_lengths, rerank_token_budget + lengths = pair_token_lengths(reranker, pairs) + slices = (None if lengths is None else + budgeted_slices(lengths, limit, rerank_token_budget())) + if slices is not None: + # Long pairs: smaller longest-first batches bound attention scratch. + order = np.array([index for group in slices for index in group], dtype=np.intp) + steps, start = {}, 0 + for group in slices: + steps[start] = len(group) + start += len(group) + recovery_name = self._reranker_name while offset < len(pairs) or scores is None: deadline.check() + step = limit if steps is None else steps[offset] indices = None if order is None else self._rerank_slice_indices( - reranker, pairs, order, offset, limit, deadline, + reranker, pairs, order, offset, step, deadline, ) - batch = (pairs[offset:offset + limit] if indices is None else + batch = (pairs[offset:offset + step] if indices is None else [pairs[index] for index in indices]) self._check_process_memory("reranker slice", deadline) deadline.check() recover = False try: - values = reranker.predict(batch, batch_size=limit, show_progress_bar=False) + values = reranker.predict(batch, batch_size=step, show_progress_bar=False) except RuntimeError as exc: from truememory.mps_utils import is_mps_oom, flush_mps_cache if not is_mps_oom(exc): @@ -1485,7 +1501,7 @@ def _handle_request_inner(self, request: dict) -> dict: if retry is not None: self._check_process_memory("reranker CPU retry", deadline) deadline.check() - values = reranker.predict(retry, batch_size=limit, show_progress_bar=False) + values = reranker.predict(retry, batch_size=len(retry) or limit, show_progress_bar=False) if indices is None: scores = _store_batch_result(scores, values, offset, len(retry), len(pairs), "scores", **result_options) diff --git a/truememory/rerank_budget.py b/truememory/rerank_budget.py new file mode 100644 index 00000000..e3f36e13 --- /dev/null +++ b/truememory/rerank_budget.py @@ -0,0 +1,85 @@ +"""Reranker input cap and token-budgeted batches that bound attention scratch.""" +from __future__ import annotations + + +# Longer pairs are truncated longest side first. Stored memories are far +# shorter; the cap bounds oversized queries. +RERANK_MAX_TOKENS = 1024 +# Attention scratch grows with batch x length^2. This budget keeps the default +# 64-pair batch for every pair up to 512 tokens. +RERANK_TOKEN_BUDGET = 64 * 512 * 512 + + +def cap_reranker_length(model: object) -> None: + """Bound a cross-encoder's input length without raising a smaller limit.""" + current = getattr(model, "max_seq_length", None) + if current is None or current > RERANK_MAX_TOKENS: + model.max_seq_length = RERANK_MAX_TOKENS + + +def rerank_token_budget(available: int | None = None) -> int: + """Per-batch attention budget, reduced while the host is short of memory.""" + if available is None: + try: + import psutil + available = psutil.virtual_memory().available + except (ImportError, OSError): + return RERANK_TOKEN_BUDGET + if available < 2 * 1024 ** 3: + return RERANK_TOKEN_BUDGET // 4 + if available < 6 * 1024 ** 3: + return RERANK_TOKEN_BUDGET // 2 + return RERANK_TOKEN_BUDGET + + +def pair_token_lengths(model: object, pairs: list) -> list[int] | None: + """Truncated token length of each pair, or None when it cannot be measured.""" + tokenizer = getattr(model, "tokenizer", None) + limit = getattr(model, "max_seq_length", None) + if (tokenizer is None or type(limit) is not int or limit <= 0 + or not all(isinstance(pair, (list, tuple)) and len(pair) == 2 + and all(isinstance(text, str) for text in pair) for pair in pairs)): + return None + if not pairs: + return [] + encoded = tokenizer([pair[0] for pair in pairs], [pair[1] for pair in pairs], + truncation="longest_first", max_length=limit) + return [len(ids) for ids in encoded["input_ids"]] + + +def budget_batch_size(length: int, limit: int, budget: int) -> int: + return max(1, min(limit, budget // max(1, length) ** 2)) + + +def budgeted_slices(lengths: list[int], limit: int, budget: int) -> list[list[int]] | None: + """Longest-first index groups sized to the budget. + + None means every default batch already fits, so callers keep their + original single call and its exact results. + """ + if budget_batch_size(max(lengths, default=0), limit, budget) >= limit: + return None + order = sorted(range(len(lengths)), key=lambda index: -lengths[index]) + slices, start = [], 0 + while start < len(order): + size = budget_batch_size(lengths[order[start]], limit, budget) + slices.append(order[start:start + size]) + start += size + return slices + + +def budgeted_predict(model: object, pairs: list, batch_size: int): + """CrossEncoder.predict with long pairs scored in smaller batches.""" + lengths = (pair_token_lengths(model, pairs) + if type(batch_size) is int and batch_size > 0 else None) + slices = None if lengths is None else budgeted_slices(lengths, batch_size, rerank_token_budget()) + if slices is None: + return model.predict(pairs, batch_size=batch_size, show_progress_bar=False) + import numpy as np + results: list = [None] * len(pairs) + for indices in slices: + values = model.predict([pairs[index] for index in indices], batch_size=len(indices), + show_progress_bar=False) + for position, index in enumerate(indices): + results[index] = values[position] + return np.stack(results) diff --git a/truememory/reranker.py b/truememory/reranker.py index cc2a6977..e4406ec0 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -32,6 +32,7 @@ from typing import TYPE_CHECKING from truememory.memory_budget import LocalMemoryPolicy +from truememory.rerank_budget import budgeted_predict, cap_reranker_length _LOCAL_MEMORY_POLICY = LocalMemoryPolicy() @@ -339,6 +340,7 @@ def get_reranker( from sentence_transformers import CrossEncoder _check(control) _model = CrossEncoder(name, device=device) + cap_reranker_length(_model) if budget_owner is not None: bind_model_owner(_model, budget_owner) except BaseException: @@ -447,7 +449,7 @@ def rerank( scope = nullcontext() with scope as memory: with _own_model(model): - scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False) + scores = budgeted_predict(model, pairs, batch_size) if memory is not None: memory.retain_buffer(scores) From ff2d6c949fb2a74e1b5a563ddcb87bbe84e65b4c Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:06:38 -0500 Subject: [PATCH 2/4] fix: Keep reranker budget helpers private and lazily imported --- tests/test-model-microbatches-735.py | 2 +- tests/test-reranker-token-budget.py | 7 ++++--- truememory/{rerank_budget.py => _rerank_budget.py} | 2 +- truememory/model_server.py | 4 ++-- truememory/reranker.py | 3 ++- 5 files changed, 10 insertions(+), 8 deletions(-) rename truememory/{rerank_budget.py => _rerank_budget.py} (98%) diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py index 0d92d236..26745d77 100644 --- a/tests/test-model-microbatches-735.py +++ b/tests/test-model-microbatches-735.py @@ -628,7 +628,7 @@ def safe_import(name: str, *args: object, **kwargs: object) -> object: return mps if name == "truememory.reranker": return types.SimpleNamespace(get_current_reranker_name=lambda: "Alibaba-NLP/gte-reranker-modernbert-base") - if name == "truememory.rerank_budget": + if name == "truememory._rerank_budget": return builtins.__import__(name, *args, **kwargs) raise AssertionError("Unexpected runtime import: " + name) diff --git a/tests/test-reranker-token-budget.py b/tests/test-reranker-token-budget.py index 62845941..1c3d1eca 100644 --- a/tests/test-reranker-token-budget.py +++ b/tests/test-reranker-token-budget.py @@ -4,8 +4,8 @@ import numpy as np import pytest -from truememory import rerank_budget -from truememory.rerank_budget import ( +from truememory import _rerank_budget as rerank_budget +from truememory._rerank_budget import ( RERANK_MAX_TOKENS, RERANK_TOKEN_BUDGET, budget_batch_size, budgeted_predict, budgeted_slices, cap_reranker_length, pair_token_lengths, rerank_token_budget, ) @@ -37,10 +37,11 @@ def _pairs(*word_counts: int) -> list[list[str]]: def test_cap_lowers_long_limits_only(): - for current, expected in ((8192, RERANK_MAX_TOKENS), (None, RERANK_MAX_TOKENS), (512, 512)): + for current, expected in ((8192, RERANK_MAX_TOKENS), (None, None), (512, 512)): model = FakeCrossEncoder(current) cap_reranker_length(model) assert model.max_seq_length == expected + cap_reranker_length(object()) def test_budget_shrinks_under_memory_pressure(): diff --git a/truememory/rerank_budget.py b/truememory/_rerank_budget.py similarity index 98% rename from truememory/rerank_budget.py rename to truememory/_rerank_budget.py index e3f36e13..c2688921 100644 --- a/truememory/rerank_budget.py +++ b/truememory/_rerank_budget.py @@ -13,7 +13,7 @@ def cap_reranker_length(model: object) -> None: """Bound a cross-encoder's input length without raising a smaller limit.""" current = getattr(model, "max_seq_length", None) - if current is None or current > RERANK_MAX_TOKENS: + if type(current) is int and current > RERANK_MAX_TOKENS: model.max_seq_length = RERANK_MAX_TOKENS diff --git a/truememory/model_server.py b/truememory/model_server.py index 8ec53bf3..7da23517 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -1078,7 +1078,7 @@ def _get_reranker( self._memory_stage("reranker construction", deadline, model_identity=(name, 0, "reranker"), device=device) model = CrossEncoder(name, device=device) - from truememory.rerank_budget import cap_reranker_length + from truememory._rerank_budget import cap_reranker_length cap_reranker_length(model) self._check_process_memory("reranker construction completion", deadline) self._reranker = model @@ -1444,7 +1444,7 @@ def _handle_request_inner(self, request: dict) -> dict: order = self._rerank_global_order(reranker, pairs, limit, deadline) if order is not None: recovery_name = self._reranker_name - from truememory.rerank_budget import budgeted_slices, pair_token_lengths, rerank_token_budget + from truememory._rerank_budget import budgeted_slices, pair_token_lengths, rerank_token_budget lengths = pair_token_lengths(reranker, pairs) slices = (None if lengths is None else budgeted_slices(lengths, limit, rerank_token_budget())) diff --git a/truememory/reranker.py b/truememory/reranker.py index e4406ec0..b12209fc 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -32,7 +32,6 @@ from typing import TYPE_CHECKING from truememory.memory_budget import LocalMemoryPolicy -from truememory.rerank_budget import budgeted_predict, cap_reranker_length _LOCAL_MEMORY_POLICY = LocalMemoryPolicy() @@ -340,6 +339,7 @@ def get_reranker( from sentence_transformers import CrossEncoder _check(control) _model = CrossEncoder(name, device=device) + from truememory._rerank_budget import cap_reranker_length cap_reranker_length(_model) if budget_owner is not None: bind_model_owner(_model, budget_owner) @@ -449,6 +449,7 @@ def rerank( scope = nullcontext() with scope as memory: with _own_model(model): + from truememory._rerank_budget import budgeted_predict scores = budgeted_predict(model, pairs, batch_size) if memory is not None: memory.retain_buffer(scores) From 40737a72af7184cf32a22744332720af39d8c9d7 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:09:17 -0500 Subject: [PATCH 3/4] fix: Touch the reranker budget only for real cross-encoders --- tests/test-model-microbatches-735.py | 2 -- truememory/model_server.py | 18 ++++++++++++------ truememory/reranker.py | 12 ++++++++---- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/tests/test-model-microbatches-735.py b/tests/test-model-microbatches-735.py index 26745d77..a241ab11 100644 --- a/tests/test-model-microbatches-735.py +++ b/tests/test-model-microbatches-735.py @@ -628,8 +628,6 @@ def safe_import(name: str, *args: object, **kwargs: object) -> object: return mps if name == "truememory.reranker": return types.SimpleNamespace(get_current_reranker_name=lambda: "Alibaba-NLP/gte-reranker-modernbert-base") - if name == "truememory._rerank_budget": - return builtins.__import__(name, *args, **kwargs) raise AssertionError("Unexpected runtime import: " + name) numpy, clock = OrderNP(), FakeClock() diff --git a/truememory/model_server.py b/truememory/model_server.py index 7da23517..9514a074 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -1078,8 +1078,9 @@ def _get_reranker( self._memory_stage("reranker construction", deadline, model_identity=(name, 0, "reranker"), device=device) model = CrossEncoder(name, device=device) - from truememory._rerank_budget import cap_reranker_length - cap_reranker_length(model) + if type(getattr(model, "max_seq_length", None)) is int: + from truememory._rerank_budget import cap_reranker_length + cap_reranker_length(model) self._check_process_memory("reranker construction completion", deadline) self._reranker = model self._reranker_name = name @@ -1444,10 +1445,15 @@ def _handle_request_inner(self, request: dict) -> dict: order = self._rerank_global_order(reranker, pairs, limit, deadline) if order is not None: recovery_name = self._reranker_name - from truememory._rerank_budget import budgeted_slices, pair_token_lengths, rerank_token_budget - lengths = pair_token_lengths(reranker, pairs) - slices = (None if lengths is None else - budgeted_slices(lengths, limit, rerank_token_budget())) + slices = None + if (getattr(reranker, "tokenizer", None) is not None + and type(getattr(reranker, "max_seq_length", None)) is int): + from truememory._rerank_budget import ( + budgeted_slices, pair_token_lengths, rerank_token_budget, + ) + lengths = pair_token_lengths(reranker, pairs) + slices = (None if lengths is None else + budgeted_slices(lengths, limit, rerank_token_budget())) if slices is not None: # Long pairs: smaller longest-first batches bound attention scratch. order = np.array([index for group in slices for index in group], dtype=np.intp) diff --git a/truememory/reranker.py b/truememory/reranker.py index b12209fc..b8099965 100644 --- a/truememory/reranker.py +++ b/truememory/reranker.py @@ -339,8 +339,9 @@ def get_reranker( from sentence_transformers import CrossEncoder _check(control) _model = CrossEncoder(name, device=device) - from truememory._rerank_budget import cap_reranker_length - cap_reranker_length(_model) + if type(getattr(_model, "max_seq_length", None)) is int: + from truememory._rerank_budget import cap_reranker_length + cap_reranker_length(_model) if budget_owner is not None: bind_model_owner(_model, budget_owner) except BaseException: @@ -449,8 +450,11 @@ def rerank( scope = nullcontext() with scope as memory: with _own_model(model): - from truememory._rerank_budget import budgeted_predict - scores = budgeted_predict(model, pairs, batch_size) + if getattr(model, "tokenizer", None) is not None and type(getattr(model, "max_seq_length", None)) is int: + from truememory._rerank_budget import budgeted_predict + scores = budgeted_predict(model, pairs, batch_size) + else: + scores = model.predict(pairs, batch_size=batch_size, show_progress_bar=False) if memory is not None: memory.retain_buffer(scores) From f1f765544fbbd3ad66fc6dc3d6f5d92f184b1ea7 Mon Sep 17 00:00:00 2001 From: buildingjoshbetter <236675935+buildingjoshbetter@users.noreply.github.com> Date: Tue, 13 Oct 2026 19:21:26 -0500 Subject: [PATCH 4/4] fix: Use the budget order directly when slicing long rerank pairs --- tests/test-reranker-token-budget.py | 15 +++++++++++++++ truememory/model_server.py | 9 ++++++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/tests/test-reranker-token-budget.py b/tests/test-reranker-token-budget.py index 1c3d1eca..90ccc6a0 100644 --- a/tests/test-reranker-token-budget.py +++ b/tests/test-reranker-token-budget.py @@ -128,3 +128,18 @@ def test_model_server_short_rerank_unchanged(monkeypatch): @pytest.mark.parametrize("length", [1, 512, 513, 1024]) def test_batch_size_never_zero_or_above_limit(length): assert 1 <= budget_batch_size(length, 64, RERANK_TOKEN_BUDGET) <= 64 + + +def test_model_server_budget_path_needs_no_input_length_helper(monkeypatch): + from truememory.model_server import ModelServer + + class WithoutInputLength(FakeCrossEncoder): + _input_length = None + + fake = WithoutInputLength() + monkeypatch.setattr(ModelServer, "_get_reranker", lambda self, model_name=None, deadline=None: fake) + monkeypatch.setattr(rerank_budget, "rerank_token_budget", lambda: RERANK_TOKEN_BUDGET) + pairs = _pairs(*([30] * 20 + [1500] * 3)) + response = ModelServer().handle_request({"op": "rerank", "pairs": pairs, "model_name": "m"}) + assert response["ok"] is True + assert np.array_equal(np.asarray(response["scores"]), FakeCrossEncoder().predict(pairs, batch_size=len(pairs))) diff --git a/truememory/model_server.py b/truememory/model_server.py index 9514a074..bed0fbdb 100644 --- a/truememory/model_server.py +++ b/truememory/model_server.py @@ -1465,9 +1465,12 @@ def _handle_request_inner(self, request: dict) -> dict: while offset < len(pairs) or scores is None: deadline.check() step = limit if steps is None else steps[offset] - indices = None if order is None else self._rerank_slice_indices( - reranker, pairs, order, offset, step, deadline, - ) + if steps is not None: + indices = order[offset:offset + step] + else: + indices = None if order is None else self._rerank_slice_indices( + reranker, pairs, order, offset, step, deadline, + ) batch = (pairs[offset:offset + step] if indices is None else [pairs[index] for index in indices]) self._check_process_memory("reranker slice", deadline)