From 0874c283f8edf223fe37b30cea74fe6eabd374ff Mon Sep 17 00:00:00 2001 From: vincent-4 Date: Sat, 15 Mar 2025 22:09:11 -0400 Subject: [PATCH 1/3] Probably fine to just LL the cross encoder re-ranker, not the image one --- lotus/models/cross_encoder_reranker.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/lotus/models/cross_encoder_reranker.py b/lotus/models/cross_encoder_reranker.py index a177ef22..6e12c549 100644 --- a/lotus/models/cross_encoder_reranker.py +++ b/lotus/models/cross_encoder_reranker.py @@ -20,7 +20,18 @@ def __init__( max_batch_size: int = 64, ): self.max_batch_size: int = max_batch_size - self.model = CrossEncoder(model, device=device) # type: ignore # CrossEncoder has wrong type stubs + self._model_name = model + self._device = device + self._model = None # Initialize model as None for lazy loading + + @property + def model(self): + """Lazy load the model when it's first accessed.""" + if self._model is None: + # Only import CrossEncoder when needed + from sentence_transformers import CrossEncoder + self._model = CrossEncoder(self._model_name, device=self._device) # type: ignore # CrossEncoder has wrong type stubs + return self._model def __call__(self, query: str, docs: list[str], K: int) -> RerankerOutput: results = self.model.rank(query, docs, top_k=K, batch_size=self.max_batch_size, show_progress_bar=False) From d26bf947d5365c7aaa7826d638060e458ec2b814 Mon Sep 17 00:00:00 2001 From: vincent-4 Date: Mon, 17 Mar 2025 17:57:55 -0400 Subject: [PATCH 2/3] ruff --- lotus/models/cross_encoder_reranker.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lotus/models/cross_encoder_reranker.py b/lotus/models/cross_encoder_reranker.py index 6e12c549..26676b71 100644 --- a/lotus/models/cross_encoder_reranker.py +++ b/lotus/models/cross_encoder_reranker.py @@ -30,6 +30,7 @@ def model(self): if self._model is None: # Only import CrossEncoder when needed from sentence_transformers import CrossEncoder + self._model = CrossEncoder(self._model_name, device=self._device) # type: ignore # CrossEncoder has wrong type stubs return self._model From b801d62401aba2ce2500d22397efbf1e883d7eaa Mon Sep 17 00:00:00 2001 From: vincent-4 Date: Mon, 17 Mar 2025 17:58:06 -0400 Subject: [PATCH 3/3] ruff fix --- lotus/models/cross_encoder_reranker.py | 1 - 1 file changed, 1 deletion(-) diff --git a/lotus/models/cross_encoder_reranker.py b/lotus/models/cross_encoder_reranker.py index 26676b71..a506660b 100644 --- a/lotus/models/cross_encoder_reranker.py +++ b/lotus/models/cross_encoder_reranker.py @@ -1,4 +1,3 @@ -from sentence_transformers import CrossEncoder from lotus.models.reranker import Reranker from lotus.types import RerankerOutput