Upload folder using huggingface_hub

README.md

@@ -120,14 +120,6 @@ To feed the data, run:
 python feed_vespa.py --vespa_app_url https://myapp.z.vespa-app.cloud --vespa_cloud_secret_token mysecrettoken
 ```
 
-### Connecting to the Vespa app and querying
-
-As a first step, you can run the `query_vespa.py` script to run some sample queries against the Vespa app:
-
-```bash
-python query_vespa.py
-```
-
 ### Starting the front-end
 
 ```bash

backend/colpali.py

@@ -13,7 +13,6 @@ import matplotlib.cm as cm
 import re
 import io
 
-import json
 import time
 import backend.testquery as testquery
 
@@ -24,12 +23,10 @@ from vidore_benchmark.interpretability.torch_utils import (
     normalize_similarity_map_per_query_token,
 )
 from vidore_benchmark.interpretability.vit_configs import VIT_CONFIG
-from vespa.application import Vespa
-from vespa.io import VespaQueryResponse
 
 matplotlib.use("Agg")
-
-MAX_QUERY_TERMS = 64
+# Prepare the colormap once to avoid recomputation
+colormap = cm.get_cmap("viridis")
 
 COLPALI_GEMMA_MODEL_NAME = "vidore/colpaligemma-3b-pt-448-base"
 
@@ -62,54 +59,6 @@ def load_vit_config(model):
     return vit_config
 
 
-def save_figure(fig, filename: str = "similarity_map.png"):
-    try:
-        OUTPUT_DIR = Path(__file__).parent.parent / "output" / "sim_maps"
-        OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
-        fig.savefig(
-            OUTPUT_DIR / filename,
-            bbox_inches="tight",
-            pad_inches=0,
-        )
-    except Exception as e:
-        print(f"Failed to save figure: {e}")
-
-
-def annotate_plot(ax, query, selected_token):
-    # Add the query text as a title over the image with opacity
-    ax.text(
-        0.5,
-        0.95,  # Adjust the position to be on the image (y=0.1 is 10% from the bottom)
-        query,
-        fontsize=18,
-        color="white",
-        ha="center",
-        va="center",
-        alpha=0.8,  # Set opacity (1 is fully opaque, 0 is fully transparent)
-        bbox=dict(
-            boxstyle="round,pad=0.5", fc="black", ec="none", lw=0, alpha=0.5
-        ),  # Add a semi-transparent background
-        transform=ax.transAxes,  # Ensure the coordinates are relative to the axes
-    )
-
-    # Add annotation with the selected token over the image with opacity
-    ax.text(
-        0.5,
-        0.05,  # Position towards the top of the image
-        f"Selected token: `{selected_token}`",
-        fontsize=18,
-        color="white",
-        ha="center",
-        va="center",
-        alpha=0.8,  # Set opacity for the text
-        bbox=dict(
-            boxstyle="round,pad=0.3", fc="black", ec="none", lw=0, alpha=0.5
-        ),  # Semi-transparent background
-        transform=ax.transAxes,  # Keep the coordinates relative to the axes
-    )
-    return ax
-
-
 def gen_similarity_maps(
     model: ColPali,
     processor: ColPaliProcessor,
@@ -140,11 +89,6 @@ def gen_similarity_maps(
 
     """
 
-    start = time.perf_counter()
-
-    # Prepare the colormap once to avoid recomputation
-    colormap = cm.get_cmap("viridis")
-
     # Process images and store original images and sizes
     processed_images = []
     original_images = []
@@ -336,154 +280,6 @@ def get_query_embeddings_and_token_map(
     return q_emb, token_to_idx
 
 
-def format_query_results(query, response, hits=5) -> dict:
-    query_time = response.json.get("timing", {}).get("searchtime", -1)
-    query_time = round(query_time, 2)
-    count = response.json.get("root", {}).get("fields", {}).get("totalCount", 0)
-    result_text = f"Query text: '{query}', query time {query_time}s, count={count}, top results:\n"
-    print(result_text)
-    return response.json
-
-
-async def query_vespa_default(
-    app: Vespa,
-    query: str,
-    q_emb: torch.Tensor,
-    hits: int = 3,
-    timeout: str = "10s",
-    **kwargs,
-) -> dict:
-    async with app.asyncio(connections=1, total_timeout=120) as session:
-        query_embedding = format_q_embs(q_emb)
-
-        start = time.perf_counter()
-        response: VespaQueryResponse = await session.query(
-            body={
-                "yql": "select id,title,url,blur_image,page_number,snippet,text,summaryfeatures from pdf_page where userQuery();",
-                "ranking": "default",
-                "query": query,
-                "timeout": timeout,
-                "hits": hits,
-                "input.query(qt)": query_embedding,
-                "presentation.timing": True,
-                **kwargs,
-            },
-        )
-        assert response.is_successful(), response.json
-        stop = time.perf_counter()
-        print(
-            f"Query time + data transfer took: {stop - start} s, vespa said searchtime was {response.json.get('timing', {}).get('searchtime', -1)} s"
-        )
-        open("response.json", "w").write(json.dumps(response.json))
-    return format_query_results(query, response)
-
-
-async def query_vespa_bm25(
-    app: Vespa,
-    query: str,
-    q_emb: torch.Tensor,
-    hits: int = 3,
-    timeout: str = "10s",
-    **kwargs,
-) -> dict:
-    async with app.asyncio(connections=1, total_timeout=120) as session:
-        query_embedding = format_q_embs(q_emb)
-
-        start = time.perf_counter()
-        response: VespaQueryResponse = await session.query(
-            body={
-                "yql": "select id,title,url,blur_image,page_number,snippet,text,summaryfeatures from pdf_page where userQuery();",
-                "ranking": "bm25",
-                "query": query,
-                "timeout": timeout,
-                "hits": hits,
-                "input.query(qt)": query_embedding,
-                "presentation.timing": True,
-                **kwargs,
-            },
-        )
-        assert response.is_successful(), response.json
-        stop = time.perf_counter()
-        print(
-            f"Query time + data transfer took: {stop - start} s, vespa said searchtime was {response.json.get('timing', {}).get('searchtime', -1)} s"
-        )
-    return format_query_results(query, response)
-
-
-def float_to_binary_embedding(float_query_embedding: dict) -> dict:
-    binary_query_embeddings = {}
-    for k, v in float_query_embedding.items():
-        binary_vector = (
-            np.packbits(np.where(np.array(v) > 0, 1, 0)).astype(np.int8).tolist()
-        )
-        binary_query_embeddings[k] = binary_vector
-        if len(binary_query_embeddings) >= MAX_QUERY_TERMS:
-            print(f"Warning: Query has more than {MAX_QUERY_TERMS} terms. Truncating.")
-            break
-    return binary_query_embeddings
-
-
-def create_nn_query_strings(
-    binary_query_embeddings: dict, target_hits_per_query_tensor: int = 20
-) -> Tuple[str, dict]:
-    # Query tensors for nearest neighbor calculations
-    nn_query_dict = {}
-    for i in range(len(binary_query_embeddings)):
-        nn_query_dict[f"input.query(rq{i})"] = binary_query_embeddings[i]
-    nn = " OR ".join(
-        [
-            f"({{targetHits:{target_hits_per_query_tensor}}}nearestNeighbor(embedding,rq{i}))"
-            for i in range(len(binary_query_embeddings))
-        ]
-    )
-    return nn, nn_query_dict
-
-
-def format_q_embs(q_embs: torch.Tensor) -> dict:
-    float_query_embedding = {k: v.tolist() for k, v in enumerate(q_embs)}
-    return float_query_embedding
-
-
-async def query_vespa_nearest_neighbor(
-    app: Vespa,
-    query: str,
-    q_emb: torch.Tensor,
-    target_hits_per_query_tensor: int = 20,
-    hits: int = 3,
-    timeout: str = "10s",
-    **kwargs,
-) -> dict:
-    # Hyperparameter for speed vs. accuracy
-    async with app.asyncio(connections=1, total_timeout=180) as session:
-        float_query_embedding = format_q_embs(q_emb)
-        binary_query_embeddings = float_to_binary_embedding(float_query_embedding)
-
-        # Mixed tensors for MaxSim calculations
-        query_tensors = {
-            "input.query(qtb)": binary_query_embeddings,
-            "input.query(qt)": float_query_embedding,
-        }
-        nn_string, nn_query_dict = create_nn_query_strings(
-            binary_query_embeddings, target_hits_per_query_tensor
-        )
-        query_tensors.update(nn_query_dict)
-        response: VespaQueryResponse = await session.query(
-            body={
-                **query_tensors,
-                "presentation.timing": True,
-                # if we use rank({nn_string}, userQuery()), dynamic summary doesn't work, see https://github.com/vespa-engine/vespa/issues/28704
-                "yql": f"select id,title,snippet,text,url,blur_image,page_number,summaryfeatures from pdf_page where {nn_string} or userQuery()",
-                "ranking.profile": "retrieval-and-rerank",
-                "timeout": timeout,
-                "hits": hits,
-                "query": query,
-                **kwargs,
-            },
-        )
-        assert response.is_successful(), response.json
-    return format_query_results(query, response)
-
-
 def is_special_token(token: str) -> bool:
     # Pattern for tokens that start with '<', numbers, whitespace, or single characters, or the string 'Question'
     # Will exclude these tokens from the similarity map generation
@@ -492,55 +288,6 @@ def is_special_token(token: str) -> bool:
         return True
     return False
 
-async def get_full_image_from_vespa(
-    app: Vespa,
-    id: str) -> str:
-    async with app.asyncio(connections=1, total_timeout=120) as session:
-        start = time.perf_counter()
-        response: VespaQueryResponse = await session.query(
-            body={
-                "yql": f"select full_image from pdf_page where id contains \"{id}\"",
-                "ranking": "unranked",
-                "presentation.timing": True,
-            },
-        )
-        assert response.is_successful(), response.json
-        stop = time.perf_counter()
-        print(
-            f"Getting image from Vespa took: {stop - start} s, vespa said searchtime was {response.json.get('timing', {}).get('searchtime', -1)} s"
-        )
-    return response.json["root"]["children"][0]["fields"]["full_image"]
-
-async def get_result_from_query(
-    app: Vespa,
-    processor: ColPaliProcessor,
-    model: ColPali,
-    query: str,
-    q_embs: torch.Tensor,
-    token_to_idx: Dict[str, int],
-    ranking: str,
-) -> Dict[str, Any]:
-    # Get the query embeddings and token map
-    print(query)
-
-    print(token_to_idx)
-    if ranking == "nn+colpali":
-        result = await query_vespa_nearest_neighbor(app, query, q_embs)
-    elif ranking == "bm25+colpali":
-        result = await query_vespa_default(app, query, q_embs)
-    elif ranking == "bm25":
-        result = await query_vespa_bm25(app, query, q_embs)
-    else:
-        raise ValueError(f"Unsupported ranking: {ranking}")
-    # Print score, title id, and text of the results
-    for idx, child in enumerate(result["root"]["children"]):
-        print(
-            f"Result {idx+1}: {child['relevance']}, {child['fields']['title']}, {child['fields']['id']}"
-        )
-    for single_result in result["root"]["children"]:
-        print(single_result["fields"].keys())
-    return result
-
 
 def add_sim_maps_to_result(
     result: Dict[str, Any],
@@ -582,36 +329,3 @@ def add_sim_maps_to_result(
     #     for token, sim_mapb64 in sim_map_dict.items():
     #         single_result["fields"][f"sim_map_{token}"] = sim_mapb64
     return result
-
-
-if __name__ == "__main__":
-    model, processor = load_model()
-    vit_config = load_vit_config(model)
-    query = "How many percent of source water is fresh water?"
-    image_filepath = (
-        Path(__file__).parent.parent
-        / "static"
-        / "assets"
-        / "ConocoPhillips Sustainability Highlights - Nature (24-0976).png"
-    )
-    q_embs, token_to_idx = get_query_embeddings_and_token_map(
-        processor,
-        model,
-        query,
-    )
-    figs_images = gen_similarity_maps(
-        model,
-        processor,
-        model.device,
-        vit_config,
-        query=query,
-        query_embs=q_embs,
-        token_idx_map=token_to_idx,
-        images=[image_filepath],
-        vespa_sim_maps=None,
-    )
-    for fig_token in figs_images:
-        for token, (fig, ax) in fig_token.items():
-            print(f"Token: {token}")
-            save_figure(fig, f"similarity_map_{token}.png")
-    print("Done")

backend/vespa_app.py

@@ -1,23 +1,332 @@
 import os
-from vespa.application import Vespa
+import time
+from typing import Dict, Any, Tuple
+
+import numpy as np
+import torch
 from dotenv import load_dotenv
+from vespa.application import Vespa
+from vespa.io import VespaQueryResponse
+
+
+class VespaQueryClient:
+    MAX_QUERY_TERMS = 64
+    VESPA_SCHEMA_NAME = "pdf_page"
+    SELECT_FIELDS = "id,title,url,blur_image,page_number,snippet,text,summaryfeatures"
 
+    def __init__(self):
+        """
+        Initialize the VespaQueryClient by loading environment variables and establishing a connection to the Vespa application.
+        """
+        load_dotenv()
+        self.vespa_app_url = os.environ.get("VESPA_APP_URL")
+        self.vespa_cloud_secret_token = os.environ.get("VESPA_CLOUD_SECRET_TOKEN")
 
-def get_vespa_app():
-    load_dotenv()
-    vespa_app_url = os.environ.get(
-        "VESPA_APP_URL"
-    )  # Ensure this is set to your Vespa app URL
-    vespa_cloud_secret_token = os.environ.get("VESPA_CLOUD_SECRET_TOKEN")
+        if not self.vespa_app_url or not self.vespa_cloud_secret_token:
+            raise ValueError(
+                "Please set the VESPA_APP_URL and VESPA_CLOUD_SECRET_TOKEN environment variables"
+            )
 
-    if not vespa_app_url or not vespa_cloud_secret_token:
-        raise ValueError(
-            "Please set the VESPA_APP_URL and VESPA_CLOUD_SECRET_TOKEN environment variables"
+        # Instantiate Vespa connection
+        self.app = Vespa(
+            url=self.vespa_app_url,
+            vespa_cloud_secret_token=self.vespa_cloud_secret_token,
         )
-    # Instantiate Vespa connection
-    vespa_app = Vespa(
-        url=vespa_app_url, vespa_cloud_secret_token=vespa_cloud_secret_token
-    )
-    vespa_app.wait_for_application_up()
-    print(f"Connected to Vespa at {vespa_app_url}")
-    return vespa_app
+        self.app.wait_for_application_up()
+        print(f"Connected to Vespa at {self.vespa_app_url}")
+
+    def format_query_results(
+        self, query: str, response: VespaQueryResponse, hits: int = 5
+    ) -> dict:
+        """
+        Format the Vespa query results.
+
+        Args:
+            query (str): The query text.
+            response (VespaQueryResponse): The response from Vespa.
+            hits (int, optional): Number of hits to display. Defaults to 5.
+
+        Returns:
+            dict: The JSON content of the response.
+        """
+        query_time = response.json.get("timing", {}).get("searchtime", -1)
+        query_time = round(query_time, 2)
+        count = response.json.get("root", {}).get("fields", {}).get("totalCount", 0)
+        result_text = f"Query text: '{query}', query time {query_time}s, count={count}, top results:\n"
+        print(result_text)
+        return response.json
+
+    async def query_vespa_default(
+        self,
+        query: str,
+        q_emb: torch.Tensor,
+        hits: int = 3,
+        timeout: str = "10s",
+        **kwargs,
+    ) -> dict:
+        """
+        Query Vespa using the default ranking profile.
+
+        Args:
+            query (str): The query text.
+            q_emb (torch.Tensor): Query embeddings.
+            hits (int, optional): Number of hits to retrieve. Defaults to 3.
+            timeout (str, optional): Query timeout. Defaults to "10s".
+
+        Returns:
+            dict: The formatted query results.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            query_embedding = self.format_q_embs(q_emb)
+
+            start = time.perf_counter()
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": (
+                        f"select {self.SELECT_FIELDS} from {self.VESPA_SCHEMA_NAME} where userQuery();"
+                    ),
+                    "ranking": "default",
+                    "query": query,
+                    "timeout": timeout,
+                    "hits": hits,
+                    "input.query(qt)": query_embedding,
+                    "presentation.timing": True,
+                    **kwargs,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            print(
+                f"Query time + data transfer took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+        return self.format_query_results(query, response)
+
+    async def query_vespa_bm25(
+        self,
+        query: str,
+        q_emb: torch.Tensor,
+        hits: int = 3,
+        timeout: str = "10s",
+        **kwargs,
+    ) -> dict:
+        """
+        Query Vespa using the BM25 ranking profile.
+
+        Args:
+            query (str): The query text.
+            q_emb (torch.Tensor): Query embeddings.
+            hits (int, optional): Number of hits to retrieve. Defaults to 3.
+            timeout (str, optional): Query timeout. Defaults to "10s".
+
+        Returns:
+            dict: The formatted query results.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            query_embedding = self.format_q_embs(q_emb)
+
+            start = time.perf_counter()
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": (
+                        f"select {self.SELECT_FIELDS} from {self.VESPA_SCHEMA_NAME} where userQuery();"
+                    ),
+                    "ranking": "bm25",
+                    "query": query,
+                    "timeout": timeout,
+                    "hits": hits,
+                    "input.query(qt)": query_embedding,
+                    "presentation.timing": True,
+                    **kwargs,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            print(
+                f"Query time + data transfer took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+        return self.format_query_results(query, response)
+
+    def float_to_binary_embedding(self, float_query_embedding: dict) -> dict:
+        """
+        Convert float query embeddings to binary embeddings.
+
+        Args:
+            float_query_embedding (dict): Dictionary of float embeddings.
+
+        Returns:
+            dict: Dictionary of binary embeddings.
+        """
+        binary_query_embeddings = {}
+        for key, vector in float_query_embedding.items():
+            binary_vector = (
+                np.packbits(np.where(np.array(vector) > 0, 1, 0))
+                .astype(np.int8)
+                .tolist()
+            )
+            binary_query_embeddings[key] = binary_vector
+            if len(binary_query_embeddings) >= self.MAX_QUERY_TERMS:
+                print(
+                    f"Warning: Query has more than {self.MAX_QUERY_TERMS} terms. Truncating."
+                )
+                break
+        return binary_query_embeddings
+
+    def create_nn_query_strings(
+        self, binary_query_embeddings: dict, target_hits_per_query_tensor: int = 20
+    ) -> Tuple[str, dict]:
+        """
+        Create nearest neighbor query strings for Vespa.
+
+        Args:
+            binary_query_embeddings (dict): Binary query embeddings.
+            target_hits_per_query_tensor (int, optional): Target hits per query tensor. Defaults to 20.
+
+        Returns:
+            Tuple[str, dict]: Nearest neighbor query string and query tensor dictionary.
+        """
+        nn_query_dict = {}
+        for i in range(len(binary_query_embeddings)):
+            nn_query_dict[f"input.query(rq{i})"] = binary_query_embeddings[i]
+        nn = " OR ".join(
+            [
+                f"({{targetHits:{target_hits_per_query_tensor}}}nearestNeighbor(embedding,rq{i}))"
+                for i in range(len(binary_query_embeddings))
+            ]
+        )
+        return nn, nn_query_dict
+
+    def format_q_embs(self, q_embs: torch.Tensor) -> dict:
+        """
+        Convert query embeddings to a dictionary of lists.
+
+        Args:
+            q_embs (torch.Tensor): Query embeddings tensor.
+
+        Returns:
+            dict: Dictionary where each key is an index and value is the embedding list.
+        """
+        return {idx: emb.tolist() for idx, emb in enumerate(q_embs)}
+
+    async def get_result_from_query(
+        self,
+        query: str,
+        q_embs: torch.Tensor,
+        ranking: str,
+        token_to_idx: dict,
+    ) -> Dict[str, Any]:
+        """
+        Get query results from Vespa based on the ranking method.
+
+        Args:
+            query (str): The query text.
+            q_embs (torch.Tensor): Query embeddings.
+            ranking (str): The ranking method to use.
+            token_to_idx (dict): Token to index mapping.
+
+        Returns:
+            Dict[str, Any]: The query results.
+        """
+        print(query)
+        print(token_to_idx)
+
+        if ranking == "nn+colpali":
+            result = await self.query_vespa_nearest_neighbor(query, q_embs)
+        elif ranking == "bm25+colpali":
+            result = await self.query_vespa_default(query, q_embs)
+        elif ranking == "bm25":
+            result = await self.query_vespa_bm25(query, q_embs)
+        else:
+            raise ValueError(f"Unsupported ranking: {ranking}")
+
+        # Print score, title id, and text of the results
+        for idx, child in enumerate(result["root"]["children"]):
+            print(
+                f"Result {idx+1}: {child['relevance']}, {child['fields']['title']}, {child['fields']['id']}"
+            )
+        for single_result in result["root"]["children"]:
+            print(single_result["fields"].keys())
+        return result
+
+    async def get_full_image_from_vespa(self, doc_id: str) -> str:
+        """
+        Retrieve the full image from Vespa for a given document ID.
+
+        Args:
+            doc_id (str): The document ID.
+
+        Returns:
+            str: The full image data.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            start = time.perf_counter()
+            response: VespaQueryResponse = await session.query(
+                body={
+                    "yql": f'select full_image from {self.VESPA_SCHEMA_NAME} where id contains "{doc_id}"',
+                    "ranking": "unranked",
+                    "presentation.timing": True,
+                },
+            )
+            assert response.is_successful(), response.json
+            stop = time.perf_counter()
+            print(
+                f"Getting image from Vespa took: {stop - start} s, Vespa reported searchtime was "
+                f"{response.json.get('timing', {}).get('searchtime', -1)} s"
+            )
+        return response.json["root"]["children"][0]["fields"]["full_image"]
+
+    async def query_vespa_nearest_neighbor(
+        self,
+        query: str,
+        q_emb: torch.Tensor,
+        target_hits_per_query_tensor: int = 20,
+        hits: int = 3,
+        timeout: str = "10s",
+        **kwargs,
+    ) -> dict:
+        """
+        Query Vespa using nearest neighbor search with mixed tensors for MaxSim calculations.
+
+        Args:
+            query (str): The query text.
+            q_emb (torch.Tensor): Query embeddings.
+            target_hits_per_query_tensor (int, optional): Target hits per query tensor. Defaults to 20.
+            hits (int, optional): Number of hits to retrieve. Defaults to 3.
+            timeout (str, optional): Query timeout. Defaults to "10s".
+
+        Returns:
+            dict: The formatted query results.
+        """
+        async with self.app.asyncio(connections=1) as session:
+            float_query_embedding = self.format_q_embs(q_emb)
+            binary_query_embeddings = self.float_to_binary_embedding(
+                float_query_embedding
+            )
+
+            # Mixed tensors for MaxSim calculations
+            query_tensors = {
+                "input.query(qtb)": binary_query_embeddings,
+                "input.query(qt)": float_query_embedding,
+            }
+            nn_string, nn_query_dict = self.create_nn_query_strings(
+                binary_query_embeddings, target_hits_per_query_tensor
+            )
+            query_tensors.update(nn_query_dict)
+
+            response: VespaQueryResponse = await session.query(
+                body={
+                    **query_tensors,
+                    "presentation.timing": True,
+                    "yql": (
+                        f"select {self.SELECT_FIELDS} from {self.VESPA_SCHEMA_NAME} where {nn_string} or userQuery()"
+                    ),
+                    "ranking.profile": "retrieval-and-rerank",
+                    "timeout": timeout,
+                    "hits": hits,
+                    "query": query,
+                    **kwargs,
+                },
+            )
+            assert response.is_successful(), response.json
+        return self.format_query_results(query, response)

main.py

@@ -13,12 +13,10 @@ from backend.cache import LRUCache
 from backend.colpali import (
     add_sim_maps_to_result,
     get_query_embeddings_and_token_map,
-    get_result_from_query,
     is_special_token,
-    get_full_image_from_vespa,
 )
 from backend.modelmanager import ModelManager
-from backend.vespa_app import get_vespa_app
+from backend.vespa_app import VespaQueryClient
 from frontend.app import (
     ChatResult,
     Home,
@@ -65,8 +63,7 @@ app, rt = fast_app(
         sselink,
     ),
 )
-vespa_app: Vespa = get_vespa_app()
-
+vespa_app: Vespa = VespaQueryClient()
 result_cache = LRUCache(max_size=20)  # Each result can be ~10MB
 task_cache = LRUCache(
     max_size=1000
@@ -173,14 +170,11 @@ async def get(request, query: str, nn: bool = True):
 
     start = time.perf_counter()
     # Fetch real search results from Vespa
-    result = await get_result_from_query(
-        app=vespa_app,
-        processor=processor,
-        model=model,
+    result = await vespa_app.get_result_from_query(
         query=query,
         q_embs=q_embs,
-        token_to_idx=token_to_idx,
         ranking=ranking_value,
+        token_to_idx=token_to_idx,
     )
     end = time.perf_counter()
     print(
@@ -278,7 +272,7 @@ async def full_image(docid: str, query_id: str, idx: int):
     """
     Endpoint to get the full quality image for a given result id.
     """
-    image_data = await get_full_image_from_vespa(vespa_app, docid)
+    image_data = await vespa_app.get_full_image_from_vespa(docid)
     # Update the cache with the full image data asynchronously to not block the request
     asyncio.create_task(update_full_image_cache(docid, query_id, idx, image_data))
     # Decode the base64 image data