import

vespa-engine · Nov 1, 2024 · 3851a52 · 3851a52
1 parent 9d203c3
commit 3851a52
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/visual-retrieval-colpali/src/backend/vespa_app.py b/visual-retrieval-colpali/src/backend/vespa_app.py
@@ -7,9 +7,10 @@
 from dotenv import load_dotenv
 from vespa.application import Vespa
 from vespa.io import VespaQueryResponse
-from .colpali import should_filter_token
+from .colpali import SimMapGenerator
 import backend.stopwords
 
+
 class VespaQueryClient:
     MAX_QUERY_TERMS = 64
     VESPA_SCHEMA_NAME = "pdf_page"
@@ -364,7 +365,7 @@ def results_to_search_results(
         fields_to_add = [
             f"sim_map_{token}_{idx}"
             for idx, token in idx_to_token.items()
-            if not should_filter_token(token)
+            if not SimMapGenerator.should_filter_token(token)
         ]
         for child in result["root"]["children"]:
             for sim_map_key in fields_to_add: