Batching

* small batch size processing for SPLADE re-ranking to work within CPU limitations
* reduce default number of context documents to 5

ask_candid/retrieval/elastic.py

@@ -299,7 +299,7 @@ def cosine_rescore(query: str, contexts: List[str]) -> List[float]:
 def reranker(
     query_results: Iterable[ElasticHitsResult],
     search_text: Optional[str] = None,
-    max_num_results: int = 10
+    max_num_results: int = 5
 ) -> Iterator[ElasticHitsResult]:
     """Reranks Elasticsearch hits coming from multiple indices/queries which may have scores on different scales.
     This will shuffle results

ask_candid/retrieval/sparse_lexical.py

@@ -1,11 +1,15 @@
 from typing import List, Dict
 
+from tqdm.auto import tqdm
+
 from transformers import AutoModelForMaskedLM, AutoTokenizer
+from torch.utils.data import DataLoader
 from torch.nn import functional as F
 import torch
 
 
 class SpladeEncoder:
+    batch_size = 4
 
     def __init__(self):
         model_id = "naver/splade-v3"
@@ -16,13 +20,16 @@ class SpladeEncoder:
 
     @torch.no_grad()
     def forward(self, texts: List[str]):
-        tokens = self.tokenizer(texts, return_tensors='pt', truncation=True, padding=True)
-        output = self.model(**tokens)
-        vec = torch.max(
-            torch.log(1 + torch.relu(output.logits)) * tokens.attention_mask.unsqueeze(-1),
-            dim=1
-        )[0].squeeze()
-        return vec
+        vectors = []
+        for batch in tqdm(DataLoader(dataset=texts, shuffle=False, batch_size=self.batch_size), desc="Re-ranking"):
+            tokens = self.tokenizer(batch, return_tensors='pt', truncation=True, padding=True)
+            output = self.model(**tokens)
+            vec = torch.max(
+                torch.log(1 + torch.relu(output.logits)) * tokens.attention_mask.unsqueeze(-1),
+                dim=1
+            )[0].squeeze()
+            vectors.append(vec)
+        return torch.vstack(vectors)
 
     def query_reranking(self, query: str, documents: List[str]):
         vec = self.forward([query, *documents])
@@ -31,7 +38,7 @@ class SpladeEncoder:
         return (xQ * xD).sum(dim=-1).cpu().tolist()
 
     def token_expand(self, query: str) -> Dict[str, float]:
-        vec = self.forward([query])
+        vec = self.forward([query]).squeeze()
         cols = vec.nonzero().squeeze().cpu().tolist()
         weights = vec[cols].cpu().tolist()