From e7da7d0e4f92ea1ae854c5d12413fd29de05e61e Mon Sep 17 00:00:00 2001 From: Fabio Scarsi <103823600+fabioscarsi@users.noreply.github.com> Date: Mon, 16 Mar 2026 15:38:33 +0100 Subject: [PATCH] feat: local reranker FP16, bucket batching, and transformers 5.x compatibility (#588) Three independent, cumulative improvements to LocalSTCrossEncoder: 1. transformers 5.x compatibility patch for XLM-RoBERTa models (Jina v2) 2. FP16 inference (opt-in via HINDSIGHT_API_RERANKER_LOCAL_FP16) 3. Length-sorted bucket batching (opt-in via HINDSIGHT_API_RERANKER_LOCAL_BUCKET_BATCHING) All behind .env switches with conservative defaults preserving current behavior. Fixes #586, Closes #587 --- hindsight-api-slim/hindsight_api/config.py | 20 ++++++ .../hindsight_api/engine/cross_encoder.py | 66 ++++++++++++++++++- 2 files changed, 84 insertions(+), 2 deletions(-) diff --git a/hindsight-api-slim/hindsight_api/config.py b/hindsight-api-slim/hindsight_api/config.py index 53d102b1..624ef679 100644 --- a/hindsight-api-slim/hindsight_api/config.py +++ b/hindsight-api-slim/hindsight_api/config.py @@ -212,6 +212,9 @@ ENV_RERANKER_LOCAL_MODEL = "HINDSIGHT_API_RERANKER_LOCAL_MODEL" ENV_RERANKER_LOCAL_FORCE_CPU = "HINDSIGHT_API_RERANKER_LOCAL_FORCE_CPU" ENV_RERANKER_LOCAL_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_LOCAL_MAX_CONCURRENT" ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE = "HINDSIGHT_API_RERANKER_LOCAL_TRUST_REMOTE_CODE" +ENV_RERANKER_LOCAL_FP16 = "HINDSIGHT_API_RERANKER_LOCAL_FP16" +ENV_RERANKER_LOCAL_BUCKET_BATCHING = "HINDSIGHT_API_RERANKER_LOCAL_BUCKET_BATCHING" +ENV_RERANKER_LOCAL_BATCH_SIZE = "HINDSIGHT_API_RERANKER_LOCAL_BATCH_SIZE" ENV_RERANKER_TEI_URL = "HINDSIGHT_API_RERANKER_TEI_URL" ENV_RERANKER_TEI_BATCH_SIZE = "HINDSIGHT_API_RERANKER_TEI_BATCH_SIZE" ENV_RERANKER_TEI_MAX_CONCURRENT = "HINDSIGHT_API_RERANKER_TEI_MAX_CONCURRENT" @@ -389,6 +392,9 @@ DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT = 4 # Limit concurrent CPU-bound rerankin DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE = ( False # Security: disabled by default, required for some models like jina-reranker-v2 ) +DEFAULT_RERANKER_LOCAL_FP16 = False # FP16 inference: opt-in, faster on MPS/CUDA (not CPU) +DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING = False # Length-sorted bucket batching: opt-in, 36-54% speedup +DEFAULT_RERANKER_LOCAL_BATCH_SIZE = 32 # Batch size for local reranker predict() calls DEFAULT_RERANKER_TEI_BATCH_SIZE = 128 DEFAULT_RERANKER_TEI_MAX_CONCURRENT = 8 DEFAULT_RERANKER_MAX_CANDIDATES = 300 @@ -668,6 +674,9 @@ class HindsightConfig: reranker_local_force_cpu: bool reranker_local_max_concurrent: int reranker_local_trust_remote_code: bool + reranker_local_fp16: bool + reranker_local_bucket_batching: bool + reranker_local_batch_size: int reranker_tei_url: str | None reranker_tei_batch_size: int reranker_tei_max_concurrent: int @@ -1092,6 +1101,17 @@ class HindsightConfig: ENV_RERANKER_LOCAL_TRUST_REMOTE_CODE, str(DEFAULT_RERANKER_LOCAL_TRUST_REMOTE_CODE) ).lower() in ("true", "1"), + reranker_local_fp16=os.getenv( + ENV_RERANKER_LOCAL_FP16, str(DEFAULT_RERANKER_LOCAL_FP16) + ).lower() + in ("true", "1"), + reranker_local_bucket_batching=os.getenv( + ENV_RERANKER_LOCAL_BUCKET_BATCHING, str(DEFAULT_RERANKER_LOCAL_BUCKET_BATCHING) + ).lower() + in ("true", "1"), + reranker_local_batch_size=int( + os.getenv(ENV_RERANKER_LOCAL_BATCH_SIZE, str(DEFAULT_RERANKER_LOCAL_BATCH_SIZE)) + ), reranker_tei_url=os.getenv(ENV_RERANKER_TEI_URL), reranker_tei_batch_size=int(os.getenv(ENV_RERANKER_TEI_BATCH_SIZE, str(DEFAULT_RERANKER_TEI_BATCH_SIZE))), reranker_tei_max_concurrent=int( diff --git a/hindsight-api-slim/hindsight_api/engine/cross_encoder.py b/hindsight-api-slim/hindsight_api/engine/cross_encoder.py index ed64c1fd..f1ac9e0b 100644 --- a/hindsight-api-slim/hindsight_api/engine/cross_encoder.py +++ b/hindsight-api-slim/hindsight_api/engine/cross_encoder.py @@ -23,6 +23,7 @@ from ..config import ( DEFAULT_RERANKER_LITELLM_MAX_TOKENS_PER_DOC, DEFAULT_RERANKER_LITELLM_MODEL, DEFAULT_RERANKER_LITELLM_SDK_MODEL, + DEFAULT_RERANKER_LOCAL_BATCH_SIZE, DEFAULT_RERANKER_LOCAL_FORCE_CPU, DEFAULT_RERANKER_LOCAL_MAX_CONCURRENT, DEFAULT_RERANKER_LOCAL_MODEL, @@ -111,6 +112,9 @@ class LocalSTCrossEncoder(CrossEncoderModel): max_concurrent: int = 4, force_cpu: bool = False, trust_remote_code: bool = False, + fp16: bool = False, + bucket_batching: bool = False, + batch_size: int = DEFAULT_RERANKER_LOCAL_BATCH_SIZE, ): """ Initialize local SentenceTransformers cross-encoder. @@ -125,10 +129,20 @@ class LocalSTCrossEncoder(CrossEncoderModel): trust_remote_code: Allow loading models with custom code (security risk). Required for some models like jina-reranker-v2-base-multilingual. Default: False (disabled for security) + fp16: Use FP16 (half precision) inference. Faster on MPS and CUDA, + may be slower on CPU. Default: False (opt-in via env var). + bucket_batching: Sort pairs by token length before batching to reduce + padding waste. 36-54% speedup, quality-identical. + Default: False (opt-in via env var). + batch_size: Batch size for predict() calls. Optimal values vary by + hardware and model (MPS: 32, CUDA: 128+). Default: 32. """ self.model_name = model_name or DEFAULT_RERANKER_LOCAL_MODEL self.force_cpu = force_cpu self.trust_remote_code = trust_remote_code + self.fp16 = fp16 + self.bucket_batching = bucket_batching + self.batch_size = batch_size self._model = None LocalSTCrossEncoder._max_concurrent = max_concurrent @@ -176,6 +190,19 @@ class LocalSTCrossEncoder(CrossEncoderModel): except Exception as e: logger.warning(f"Failed to detect GPU/MPS, falling back to CPU: {e}") + # Patch transformers 5.x compatibility for models using XLM-RoBERTa + # (e.g., jina-reranker-v2-base-multilingual). transformers 5.x removed + # create_position_ids_from_input_ids as a module-level function; the custom + # code in these models still references it. This monkey-patch restores it. + try: + from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaEmbeddings + import transformers.models.xlm_roberta.modeling_xlm_roberta as xlm_module + if not hasattr(xlm_module, 'create_position_ids_from_input_ids'): + xlm_module.create_position_ids_from_input_ids = XLMRobertaEmbeddings.create_position_ids_from_input_ids + logger.info("Reranker: applied transformers 5.x compatibility patch for XLM-RoBERTa") + except Exception: + pass + # Suppress verbose transformers warnings during model loading # This suppresses the "UNEXPECTED" warnings from CrossEncoder which are harmless # but look alarming to users (e.g., "embeddings.position_ids | UNEXPECTED") @@ -200,6 +227,12 @@ class LocalSTCrossEncoder(CrossEncoderModel): # Restore original logging level transformers_logger.setLevel(original_level) + # FP16 inference: convert model weights to half precision. + # Empirically validated: 27-36% faster on MPS, quality-identical (20/20 overlap). + if self.fp16 and device != "cpu": + self._model.model.half() + logger.info("Reranker: FP16 inference enabled") + # Initialize shared executor (limited workers naturally limits concurrency) if LocalSTCrossEncoder._executor is None: LocalSTCrossEncoder._executor = ThreadPoolExecutor( @@ -211,8 +244,34 @@ class LocalSTCrossEncoder(CrossEncoderModel): logger.info("Reranker: local provider initialized (using existing executor)") def _predict_sync(self, pairs: list[tuple[str, str]]) -> list[float]: - """Synchronous prediction wrapper for thread pool execution.""" - scores = self._model.predict(pairs, show_progress_bar=False) + """Synchronous prediction wrapper for thread pool execution. + + Supports two optimizations (controlled via .env): + - bucket_batching: sort pairs by token length to reduce padding waste (36-54% speedup) + - batch_size: explicit batch size for predict() calls (MPS optimal: 32) + """ + import numpy as np + + if self.bucket_batching and len(pairs) > 1: + # Sort pairs by approximate token length to create homogeneous batches. + # This eliminates padding waste — short pairs aren't padded to the length + # of the longest pair in the batch. Quality-identical by construction. + lengths = [len(pairs[i][0]) + len(pairs[i][1]) for i in range(len(pairs))] + sorted_indices = sorted(range(len(pairs)), key=lambda i: lengths[i]) + sorted_pairs = [pairs[i] for i in sorted_indices] + + sorted_scores = self._model.predict( + sorted_pairs, batch_size=self.batch_size, show_progress_bar=False + ) + sorted_scores = sorted_scores.tolist() if hasattr(sorted_scores, "tolist") else list(sorted_scores) + + # Restore original order + scores = [0.0] * len(pairs) + for new_pos, orig_idx in enumerate(sorted_indices): + scores[orig_idx] = sorted_scores[new_pos] + return scores + + scores = self._model.predict(pairs, batch_size=self.batch_size, show_progress_bar=False) return scores.tolist() if hasattr(scores, "tolist") else list(scores) async def predict(self, pairs: list[tuple[str, str]]) -> list[float]: @@ -1196,6 +1255,9 @@ def create_cross_encoder_from_env() -> CrossEncoderModel: max_concurrent=config.reranker_local_max_concurrent, force_cpu=config.reranker_local_force_cpu, trust_remote_code=config.reranker_local_trust_remote_code, + fp16=config.reranker_local_fp16, + bucket_batching=config.reranker_local_bucket_batching, + batch_size=config.reranker_local_batch_size, ) elif provider == "cohere": api_key = config.reranker_cohere_api_key