fix bug (castorini#793)

fix key not found in idf dict
MXueguang · Nov 5, 2021 · fe55b35 · fe55b35
1 parent d47f63f
commit fe55b35
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pyserini/search/_impact_searcher.py b/pyserini/search/_impact_searcher.py
@@ -112,7 +112,7 @@ def search(self, q: str, k: int = 10, fields=dict()) -> List[JImpactSearcherResu
         encoded_query = self.query_encoder.encode(q)
         jquery = JHashMap()
         for (token, weight) in encoded_query.items():
-            if self.idf[token] > self.min_idf:
+            if token in self.idf and self.idf[token] > self.min_idf:
                 jquery.put(JString(token.encode('utf8')), JFloat(weight))
 
         if not fields:
@@ -153,7 +153,7 @@ def batch_search(self, queries: List[str], qids: List[str],
             encoded_query = self.query_encoder.encode(q)
             jquery = JHashMap()
             for (token, weight) in encoded_query.items():
-                if self.idf[token] > self.min_idf:
+                if token in self.idf and self.idf[token] > self.min_idf:
                     jquery.put(JString(token.encode('utf8')), JFloat(weight))
             query_lst.add(jquery)