From 3851a5230c8b4706d6c33311981fbe93ec96fc82 Mon Sep 17 00:00:00 2001
From: thomasht86 <thomas@vespa.ai>
Date: Fri, 1 Nov 2024 15:38:14 +0100
Subject: [PATCH] import

---
 visual-retrieval-colpali/src/backend/vespa_app.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/visual-retrieval-colpali/src/backend/vespa_app.py b/visual-retrieval-colpali/src/backend/vespa_app.py
index c553a3d8f..b5a53c76d 100644
--- 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: