Update splade_encoder.py

splade_encoder.py

@@ -50,7 +50,6 @@ class SpladeEmbeddingFunction(BaseEmbeddingFunction):
     def __init__(
         self,
         model_name: str = "naver/splade-cocondenser-ensembledistil",
-        batch_size: int = 32,
         query_instruction: str = "",
         doc_instruction: str = "",
         device: Optional[str] = "cpu",
@@ -60,7 +59,7 @@ class SpladeEmbeddingFunction(BaseEmbeddingFunction):
         self.model_name = model_name
 
         _model_config = dict(
-            {"model_name_or_path": model_name, "batch_size": batch_size, "device": device}
+            {"model_name_or_path": model_name, "device": device}
         )
         self._model_config = _model_config
         self.model = _SpladeImplementation(**self._model_config)
@@ -70,7 +69,7 @@ class SpladeEmbeddingFunction(BaseEmbeddingFunction):
         self.query_instruction = query_instruction
         self.doc_instruction = doc_instruction
 
-    def __call__(self, texts: List[str]) -> csr_array:
+    def __call__(self, texts: List[str], batch_size: int = 32) -> csr_array:
         return self._encode(texts, None)
 
     def encode_documents(self, documents: List[str]) -> csr_array:
@@ -78,8 +77,8 @@ class SpladeEmbeddingFunction(BaseEmbeddingFunction):
             [self.doc_instruction + document for document in documents], self.k_tokens_document,
         )
 
-    def _encode(self, texts: List[str], k_tokens: int) -> csr_array:
-        return self.model.forward(texts, k_tokens=k_tokens)
+    def _encode(self, texts: List[str], k_tokens: int, batch_size: int) -> csr_array:
+        return self.model.forward(texts, k_tokens=k_tokens, batch_size=batch_size)
 
     def encode_queries(self, queries: List[str]) -> csr_array:
         return self._encode(
@@ -103,8 +102,7 @@ class _SpladeImplementation:
     def __init__(
         self,
         model_name_or_path: Optional[str] = None,
-        device: Optional[str] = None,
-        batch_size: int = 32
+        device: Optional[str] = None
     ):
         self.device = device
         self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
@@ -119,7 +117,6 @@ class _SpladeImplementation:
             use_io_binding=True,
             session_options=session_options
         )
-        self.batch_size = batch_size
 
         self.relu = torch.nn.ReLU()
         self.relu.to(self.device)
@@ -141,9 +138,9 @@ class _SpladeImplementation:
     def _batchify(self, texts: List[str], batch_size: int) -> List[List[str]]:
         return [texts[i : i + batch_size] for i in range(0, len(texts), batch_size)]
 
-    def forward(self, texts: List[str], k_tokens: int) -> csr_array:
+    def forward(self, texts: List[str], k_tokens: int, batch_size: int) -> csr_array:
         with torch.no_grad():
-            batched_texts = self._batchify(texts, self.batch_size)
+            batched_texts = self._batchify(texts, batch_size)
             sparse_embs = []
             for batch_texts in batched_texts:
                 logits = self._encode(texts=batch_texts)