Upload folder using huggingface_hub

README.md

@@ -102,8 +102,14 @@ python feed_vespa.py --vespa_app_url https://myapp.z.vespa-app.cloud --vespa_clo
 
 ### Connecting to the Vespa app and querying
 
-As a first step, until we hook up to frontend, you can run the `query_vespa.py` script to run some sample queries against the Vespa app:
+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
+python main.py
+```

backend/cache.py

@@ -0,0 +1,26 @@
+from collections import OrderedDict
+
+
+# Initialize LRU Cache
+class LRUCache:
+    def __init__(self, max_size=20):
+        self.max_size = max_size
+        self.cache = OrderedDict()
+
+    def get(self, key):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+            return self.cache[key]
+        return None
+
+    def set(self, key, value):
+        if key in self.cache:
+            self.cache.move_to_end(key)
+        else:
+            if len(self.cache) >= self.max_size:
+                self.cache.popitem(last=False)
+        self.cache[key] = value
+
+    def delete(self, key):
+        if key in self.cache:
+            del self.cache[key]

backend/colpali.py

@@ -4,45 +4,33 @@ import torch
 from PIL import Image
 import numpy as np
 from typing import cast
-import pprint
 from pathlib import Path
 import base64
 from io import BytesIO
-from typing import Union, Tuple
+from typing import Union, Tuple, List, Dict, Any
 import matplotlib
+import matplotlib.cm as cm
 import re
+import io
+
+import json
+import time
 
 from colpali_engine.models import ColPali, ColPaliProcessor
 from colpali_engine.utils.torch_utils import get_torch_device
 from einops import rearrange
-from vidore_benchmark.interpretability.plot_utils import plot_similarity_heatmap
 from vidore_benchmark.interpretability.torch_utils import (
     normalize_similarity_map_per_query_token,
 )
 from vidore_benchmark.interpretability.vit_configs import VIT_CONFIG
-from vidore_benchmark.utils.image_utils import scale_image
 from vespa.application import Vespa
 from vespa.io import VespaQueryResponse
 
 matplotlib.use("Agg")
 
 MAX_QUERY_TERMS = 64
-# OUTPUT_DIR = Path(__file__).parent.parent / "output" / "sim_maps"
-# OUTPUT_DIR.mkdir(exist_ok=True)
-
-COLPALI_GEMMA_MODEL_ID = "vidore--colpaligemma-3b-pt-448-base"
-COLPALI_GEMMA_MODEL_SNAPSHOT = "12c59eb7e23bc4c26876f7be7c17760d5d3a1ffa"
-COLPALI_GEMMA_MODEL_PATH = (
-    Path().home()
-    / f".cache/huggingface/hub/models--{COLPALI_GEMMA_MODEL_ID}/snapshots/{COLPALI_GEMMA_MODEL_SNAPSHOT}"
-)
-COLPALI_MODEL_ID = "vidore--colpali-v1.2"
-COLPALI_MODEL_SNAPSHOT = "9912ce6f8a462d8cf2269f5606eabbd2784e764f"
-COLPALI_MODEL_PATH = (
-    Path().home()
-    / f".cache/huggingface/hub/models--{COLPALI_MODEL_ID}/snapshots/{COLPALI_MODEL_SNAPSHOT}"
-)
-COLPALI_GEMMA_MODEL_NAME = COLPALI_GEMMA_MODEL_ID.replace("--", "/")
+
+COLPALI_GEMMA_MODEL_NAME = "vidore/colpaligemma-3b-pt-448-base"
 
 
 def load_model() -> Tuple[ColPali, ColPaliProcessor]:
@@ -73,195 +61,241 @@ def load_vit_config(model):
     return vit_config
 
 
-# Create dummy image
-dummy_image = Image.new("RGB", (448, 448), (255, 255, 255))
-
-
-def gen_similarity_map(
-    model, processor, device, vit_config, query, image: Union[Path, str]
-):
-    # Should take in the b64 image from Vespa query result
-    # And possibly the tensor representing the output_image
-    if isinstance(image, Path):
-        # image is a file path
-        try:
-            image = Image.open(image)
-        except Exception as e:
-            raise ValueError(f"Failed to open image from path: {e}")
-    elif isinstance(image, str):
-        # image is b64 string
-        try:
-            image = Image.open(BytesIO(base64.b64decode(image)))
-        except Exception as e:
-            raise ValueError(f"Failed to open image from b64: {e}")
-
-    # Preview the image
-    scale_image(image, 512)
-    # Preprocess inputs
-    input_text_processed = processor.process_queries([query]).to(device)
-    input_image_processed = processor.process_images([image]).to(device)
-    # Forward passes
-    with torch.no_grad():
-        output_text = model.forward(**input_text_processed)
-        output_image = model.forward(**input_image_processed)
-    # output_image is the tensor that we could get from the Vespa query
-    # Print shape of output_text and output_image
-    # Output image shape: torch.Size([1, 1030, 128])
-    # Remove the special tokens from the output
-    output_image = output_image[
-        :, : processor.image_seq_length, :
-    ]  # (1, n_patches_x * n_patches_y, dim)
-
-    # Rearrange the output image tensor to explicitly represent the 2D grid of patches
-    output_image = rearrange(
-        output_image,
-        "b (h w) c -> b h w c",
-        h=vit_config.n_patch_per_dim,
-        w=vit_config.n_patch_per_dim,
-    )  # (1, n_patches_x, n_patches_y, dim)
-    # Get the similarity map
-    similarity_map = torch.einsum(
-        "bnk,bijk->bnij", output_text, output_image
-    )  # (1, query_tokens, n_patches_x, n_patches_y)
-
-    # Normalize the similarity map
-    similarity_map_normalized = normalize_similarity_map_per_query_token(
-        similarity_map
-    )  # (1, query_tokens, n_patches_x, n_patches_y)
-    # Use this cell output to choose a token using its index
-    query_tokens = processor.tokenizer.tokenize(
-        processor.decode(input_text_processed.input_ids[0])
-    )
-    # Choose a token
-    token_idx = (
-        10  # e.g. if "12: '▁Kazakhstan',", set 12 to choose the token 'Kazakhstan'
-    )
-    selected_token = processor.decode(input_text_processed.input_ids[0, token_idx])
-    # strip whitespace
-    selected_token = selected_token.strip()
-    print(f"Selected token: `{selected_token}`")
-    # Retrieve the similarity map for the chosen token
-    pprint.pprint({idx: val for idx, val in enumerate(query_tokens)})
-    # Resize the image to square
-    input_image_square = image.resize((vit_config.resolution, vit_config.resolution))
-
-    # Plot the similarity map
-    fig, ax = plot_similarity_heatmap(
-        input_image_square,
-        patch_size=vit_config.patch_size,
-        image_resolution=vit_config.resolution,
-        similarity_map=similarity_map_normalized[0, token_idx, :, :],
-    )
-    ax = annotate_plot(ax, selected_token)
-    return fig, ax
-
-
-# def save_figure(fig, filename: str = "similarity_map.png"):
-#     fig.savefig(
-#         OUTPUT_DIR / filename,
-#         bbox_inches="tight",
-#         pad_inches=0,
-#     )
+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
-    ax.set_title(query, fontsize=18)
-    # Add annotation with selected token
-    ax.annotate(
-        f"Selected token:`{selected_token}`",
-        xy=(0.5, 0.95),
-        xycoords="axes fraction",
+    # 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="black",
-        bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="black", lw=1),
+        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_map_new(
-    processor: ColPaliProcessor,
+def gen_similarity_maps(
     model: ColPali,
+    processor: ColPaliProcessor,
     device,
     vit_config,
     query: str,
     query_embs: torch.Tensor,
     token_idx_map: dict,
-    token_to_show: str,
-    image: Union[Path, str],
-):
-    if isinstance(image, Path):
-        # image is a file path
-        try:
-            image = Image.open(image)
-        except Exception as e:
-            raise ValueError(f"Failed to open image from path: {e}")
-    elif isinstance(image, str):
-        # image is b64 string
-        try:
-            image = Image.open(BytesIO(base64.b64decode(image)))
-        except Exception as e:
-            raise ValueError(f"Failed to open image from b64: {e}")
-    token_idx = token_idx_map[token_to_show]
-    print(f"Selected token: `{token_to_show}`")
-    # strip whitespace
-    # Preview the image
-    # scale_image(image, 512)
-    # Preprocess inputs
-    input_image_processed = processor.process_images([image]).to(device)
-    # Forward passes
-    with torch.no_grad():
-        output_image = model.forward(**input_image_processed)
-    # output_image is the tensor that we could get from the Vespa query
-    # Print shape of output_text and output_image
-    # Output image shape: torch.Size([1, 1030, 128])
-    # Remove the special tokens from the output
-    print(f"Output image shape before dim: {output_image.shape}")
-    output_image = output_image[
-        :, : processor.image_seq_length, :
-    ]  # (1, n_patches_x * n_patches_y, dim)
-    print(f"Output image shape after dim: {output_image.shape}")
-    # Rearrange the output image tensor to explicitly represent the 2D grid of patches
-    output_image = rearrange(
-        output_image,
-        "b (h w) c -> b h w c",
-        h=vit_config.n_patch_per_dim,
-        w=vit_config.n_patch_per_dim,
-    )  # (1, n_patches_x, n_patches_y, dim)
-    # Get the similarity map
-    print(f"Query embs shape: {query_embs.shape}")
-    # Add 1 extra dim to start of query_embs
-    query_embs = query_embs.unsqueeze(0).to(device)
-    print(f"Output image shape: {output_image.shape}")
-    similarity_map = torch.einsum(
-        "bnk,bijk->bnij", query_embs, output_image
-    )  # (1, query_tokens, n_patches_x, n_patches_y)
-    print(f"Similarity map shape: {similarity_map.shape}")
-    # Normalize the similarity map
-    similarity_map_normalized = normalize_similarity_map_per_query_token(
-        similarity_map
-    )  # (1, query_tokens, n_patches_x, n_patches_y)
-    print(f"Similarity map normalized shape: {similarity_map_normalized.shape}")
-    # Use this cell output to choose a token using its index
-    input_image_square = image.resize((vit_config.resolution, vit_config.resolution))
-
-    # Plot the similarity map
-    fig, ax = plot_similarity_heatmap(
-        input_image_square,
-        patch_size=vit_config.patch_size,
-        image_resolution=vit_config.resolution,
-        similarity_map=similarity_map_normalized[0, token_idx, :, :],
-    )
-    ax = annotate_plot(ax, query, token_to_show)
-    # save the figure
-    # save_figure(fig, f"similarity_map_{token_to_show}.png")
-    return fig, ax
+    images: List[Union[Path, str]],
+    vespa_sim_maps: List[str],
+) -> List[Dict[str, str]]:
+    """
+    Generate similarity maps for the given images and query, and return base64-encoded blended images.
+
+    Args:
+        model (ColPali): The model used for generating embeddings.
+        processor (ColPaliProcessor): Processor for images and text.
+        device: Device to run the computations on.
+        vit_config: Configuration for the Vision Transformer.
+        query (str): The query string.
+        query_embs (torch.Tensor): Query embeddings.
+        token_idx_map (dict): Mapping from tokens to their indices.
+        images (List[Union[Path, str]]): List of image paths or base64-encoded strings.
+
+    Returns:
+        List[Dict[str, str]]: A list where each item is a dictionary mapping tokens to base64-encoded blended images.
+    """
+
+    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 = []
+    original_sizes = []
+    for img in images:
+        if isinstance(img, Path):
+            try:
+                img_pil = Image.open(img).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"Failed to open image from path: {e}")
+        elif isinstance(img, str):
+            try:
+                img_pil = Image.open(BytesIO(base64.b64decode(img))).convert("RGB")
+            except Exception as e:
+                raise ValueError(f"Failed to open image from base64 string: {e}")
+        else:
+            raise ValueError(f"Unsupported image type: {type(img)}")
+        original_images.append(img_pil.copy())
+        original_sizes.append(img_pil.size)  # (width, height)
+        processed_images.append(img_pil)
+
+    # If similarity maps are provided, use them instead of computing them
+    if vespa_sim_maps:
+        print("Using provided similarity maps")
+        # A sim map looks like this:
+        # "similarities": [
+        #      {
+        #        "address": {
+        #          "patch": "0",
+        #          "querytoken": "0"
+        #        },
+        #        "value": 1.2599412202835083
+        #      },
+        # ... and so on.
+        # Now turn these into a tensor of same shape as previous similarity map
+        vespa_sim_map_tensor = torch.zeros(
+                (len(vespa_sim_maps), query_embs.size(dim=1), vit_config.n_patch_per_dim, vit_config.n_patch_per_dim)
+            )
+        for idx, vespa_sim_map in enumerate(vespa_sim_maps):
+            for cell in vespa_sim_map["similarities"]["cells"]:
+                patch = int(cell["address"]["patch"])
+                if patch >= processor.image_seq_length:
+                    continue
+                query_token = int(cell["address"]["querytoken"])
+                value = cell["value"]
+                vespa_sim_map_tensor[idx, int(query_token), int(patch) // vit_config.n_patch_per_dim, int(patch) % vit_config.n_patch_per_dim] = value            
+
+        # Normalize the similarity map per query token
+        similarity_map_normalized = normalize_similarity_map_per_query_token(vespa_sim_map_tensor)
+    else:
+        # Preprocess inputs
+        print("Computing similarity maps")
+        start2 = time.perf_counter()
+        input_image_processed = processor.process_images(processed_images).to(device)
+
+        # Forward passes
+        with torch.no_grad():
+            output_image = model.forward(**input_image_processed)
+
+        # Remove the special tokens from the output
+        output_image = output_image[:, : processor.image_seq_length, :]
+
+        # Rearrange the output image tensor to represent the 2D grid of patches
+        output_image = rearrange(
+            output_image,
+            "b (h w) c -> b h w c",
+            h=vit_config.n_patch_per_dim,
+            w=vit_config.n_patch_per_dim,
+        )
+
+        # Ensure query_embs has batch dimension
+        if query_embs.dim() == 2:
+            query_embs = query_embs.unsqueeze(0).to(device)
+        else:
+            query_embs = query_embs.to(device)
+
+        # Compute the similarity map
+        similarity_map = torch.einsum(
+            "bnk,bhwk->bnhw", query_embs, output_image
+        )  # Shape: (batch_size, query_tokens, h, w)
+
+        end2 = time.perf_counter()
+        print(f"Similarity map computation took: {end2 - start2} s")
+
+        # Normalize the similarity map per query token
+        similarity_map_normalized = normalize_similarity_map_per_query_token(similarity_map)
+
+    # Collect the blended images
+    start3 = time.perf_counter()
+    results = []
+    for idx, img in enumerate(original_images):
+        original_size = original_sizes[idx]  # (width, height)
+        result_per_image = {}
+        for token, token_idx in token_idx_map.items():
+            if is_special_token(token):
+                continue
+
+            # Get the similarity map for this image and the selected token
+            sim_map = similarity_map_normalized[idx, token_idx, :, :]  # Shape: (h, w)
+
+            # Move the similarity map to CPU and convert to NumPy array
+            sim_map_np = sim_map.cpu().numpy()
+
+            # Resize the similarity map to the original image size
+            sim_map_img = Image.fromarray(sim_map_np)
+            sim_map_resized = sim_map_img.resize(original_size, resample=Image.BICUBIC)
+
+            # Convert the resized similarity map to a NumPy array
+            sim_map_resized_np = np.array(sim_map_resized, dtype=np.float32)
+
+            # Normalize the similarity map to range [0, 1]
+            sim_map_min = sim_map_resized_np.min()
+            sim_map_max = sim_map_resized_np.max()
+            if sim_map_max - sim_map_min > 1e-6:
+                sim_map_normalized = (sim_map_resized_np - sim_map_min) / (
+                    sim_map_max - sim_map_min
+                )
+            else:
+                sim_map_normalized = np.zeros_like(sim_map_resized_np)
+
+            # Apply a colormap to the normalized similarity map
+            heatmap = colormap(sim_map_normalized)  # Returns an RGBA array
+
+            # Convert the heatmap to a PIL Image
+            heatmap_uint8 = (heatmap * 255).astype(np.uint8)
+            heatmap_img = Image.fromarray(heatmap_uint8)
+
+            # Ensure both images are in RGBA mode
+            original_img_rgba = img.convert("RGBA")
+            heatmap_img_rgba = heatmap_img.convert("RGBA")
+
+            # Overlay the heatmap onto the original image
+            blended_img = Image.blend(
+                original_img_rgba, heatmap_img_rgba, alpha=0.4
+            )  # Adjust alpha as needed
+            # Save the blended image to a BytesIO buffer
+            buffer = io.BytesIO()
+            blended_img.save(buffer, format="PNG")
+            buffer.seek(0)
+
+            # Encode the image to base64
+            blended_img_base64 = base64.b64encode(buffer.read()).decode("utf-8")
+
+            # Store the base64-encoded image
+            result_per_image[token] = blended_img_base64
+        results.append(result_per_image)
+    end3 = time.perf_counter()
+    print(f"Collecting blended images took: {end3 - start3} s")
+    print(f"Total heatmap generation took: {end3 - start} s")
+    return results
 
 
 def get_query_embeddings_and_token_map(
-    processor, model, query, image
+    processor, model, query
 ) -> Tuple[torch.Tensor, dict]:
     inputs = processor.process_queries([query]).to(model.device)
     with torch.no_grad():
@@ -294,9 +328,11 @@ async def query_vespa_default(
 ) -> 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,image,page_number,text from pdf_page where userQuery();",
+                "yql": "select id,title,url,full_image,page_number,snippet,text,summaryfeatures from pdf_page where userQuery();",
                 "ranking": "default",
                 "query": query,
                 "timeout": timeout,
@@ -307,6 +343,32 @@ async def query_vespa_default(
             },
         )
         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,
+    hits: int = 3,
+    timeout: str = "10s",
+    **kwargs,
+) -> dict:
+    async with app.asyncio(connections=1, total_timeout=120) as session:
+        response: VespaQueryResponse = await session.query(
+            body={
+                "yql": "select id,title,url,full_image,page_number,snippet,text from pdf_page where userQuery();",
+                "ranking": "bm25",
+                "query": query,
+                "timeout": timeout,
+                "hits": hits,
+                "presentation.timing": True,
+                **kwargs,
+            },
+        )
+        assert response.is_successful(), response.json
     return format_query_results(query, response)
 
 
@@ -371,10 +433,12 @@ async def query_vespa_nearest_neighbor(
             body={
                 **query_tensors,
                 "presentation.timing": True,
-                "yql": f"select id,title,text,url,image,page_number from pdf_page where {nn_string}",
+                # 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,full_image,page_number from pdf_page where {nn_string} or userQuery()",
                 "ranking.profile": "retrieval-and-rerank",
                 "timeout": timeout,
                 "hits": hits,
+                "query": query,
                 **kwargs,
             },
         )
@@ -383,8 +447,8 @@ async def query_vespa_nearest_neighbor(
 
 
 def is_special_token(token: str) -> bool:
-    # Pattern for tokens that start with '<', numbers, whitespace, or single characters
-    pattern = re.compile(r"^<.*$|^\d+$|^\s+$|^.$")
+    # Pattern for tokens that start with '<', numbers, whitespace, or single characters, or the string 'Question'
+    pattern = re.compile(r"^<.*$|^\d+$|^\s+$|^\w$|^Question$")
     if pattern.match(token):
         return True
     return False
@@ -395,111 +459,64 @@ async def get_result_from_query(
     processor: ColPaliProcessor,
     model: ColPali,
     query: str,
-    nn=False,
-    gen_sim_map=False,
-):
+    q_embs: torch.Tensor,
+    token_to_idx: Dict[str, int],
+    ranking: str,
+) -> Dict[str, Any]:
     # Get the query embeddings and token map
     print(query)
-    q_embs, token_to_idx = get_query_embeddings_and_token_map(
-        processor, model, query, dummy_image
-    )
-    print(token_to_idx)
-    # Use the token map to choose a token randomly for now
-    # Dynamically select a token containing 'water'
 
-    if nn:
+    print(token_to_idx)
+    if ranking == "nn+colpali":
         result = await query_vespa_nearest_neighbor(app, query, q_embs)
-    else:
+    elif ranking == "bm25+colpali":
         result = await query_vespa_default(app, query, q_embs)
-    # Print score, title id and text of the results
+    elif ranking == "bm25":
+        result = await query_vespa_bm25(app, query)
+    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']}"
         )
-
-    if gen_sim_map:
-        for single_result in result["root"]["children"]:
-            img = single_result["fields"]["image"]
-            for token in token_to_idx:
-                if is_special_token(token):
-                    print(f"Skipping special token: {token}")
-                    continue
-                fig, ax = gen_similarity_map_new(
-                    processor,
-                    model,
-                    model.device,
-                    load_vit_config(model),
-                    query,
-                    q_embs,
-                    token_to_idx,
-                    token,
-                    img,
-                )
-                sim_map = base64.b64encode(fig.canvas.tostring_rgb()).decode("utf-8")
-                single_result["fields"][f"sim_map_{token}"] = sim_map
+    for single_result in result["root"]["children"]:
+        print(single_result["fields"].keys())
     return result
 
 
-def get_result_dummy(query: str, nn: bool = False):
-    result = {}
-    result["timing"] = {}
-    result["timing"]["querytime"] = 0.23700000000000002
-    result["timing"]["summaryfetchtime"] = 0.001
-    result["timing"]["searchtime"] = 0.23900000000000002
-    result["root"] = {}
-    result["root"]["id"] = "toplevel"
-    result["root"]["relevance"] = 1
-    result["root"]["fields"] = {}
-    result["root"]["fields"]["totalCount"] = 59
-    result["root"]["coverage"] = {}
-    result["root"]["coverage"]["coverage"] = 100
-    result["root"]["coverage"]["documents"] = 155
-    result["root"]["coverage"]["full"] = True
-    result["root"]["coverage"]["nodes"] = 1
-    result["root"]["coverage"]["results"] = 1
-    result["root"]["coverage"]["resultsFull"] = 1
-    result["root"]["children"] = []
-    elt0 = {}
-    elt0["id"] = "index:colpalidemo_content/0/424c85e7dece761d226f060f"
-    elt0["relevance"] = 2354.050122871995
-    elt0["source"] = "colpalidemo_content"
-    elt0["fields"] = {}
-    elt0["fields"]["id"] = "a767cb1868be9a776cd56b768347b089"
-    elt0["fields"]["url"] = (
-        "https://static.conocophillips.com/files/resources/conocophillips-2023-sustainability-report.pdf"
-    )
-    elt0["fields"]["title"] = "ConocoPhillips 2023 Sustainability Report"
-    elt0["fields"]["page_number"] = 50
-    elt0["fields"]["image"] = "empty for now - is base64 encoded image"
-    result["root"]["children"].append(elt0)
-    elt1 = {}
-    elt1["id"] = "index:colpalidemo_content/0/b927c4979f0beaf0d7fab8e9"
-    elt1["relevance"] = 2313.7529950886965
-    elt1["source"] = "colpalidemo_content"
-    elt1["fields"] = {}
-    elt1["fields"]["id"] = "9f2fc0aa02c9561adfaa1451c875658f"
-    elt1["fields"]["url"] = (
-        "https://static.conocophillips.com/files/resources/conocophillips-2023-managing-climate-related-risks.pdf"
-    )
-    elt1["fields"]["title"] = "ConocoPhillips Managing Climate Related Risks"
-    elt1["fields"]["page_number"] = 44
-    elt1["fields"]["image"] = "empty for now - is base64 encoded image"
-    result["root"]["children"].append(elt1)
-    elt2 = {}
-    elt2["id"] = "index:colpalidemo_content/0/9632d72238829d6afefba6c9"
-    elt2["relevance"] = 2312.230182081461
-    elt2["source"] = "colpalidemo_content"
-    elt2["fields"] = {}
-    elt2["fields"]["id"] = "d638ded1ddcb446268b289b3f65430fd"
-    elt2["fields"]["url"] = (
-        "https://static.conocophillips.com/files/resources/24-0976-sustainability-highlights_nature.pdf"
-    )
-    elt2["fields"]["title"] = (
-        "ConocoPhillips Sustainability Highlights - Nature (24-0976)"
+def add_sim_maps_to_result(
+    result: Dict[str, Any],
+    model: ColPali,
+    processor: ColPaliProcessor,
+    query: str,
+    q_embs: Any,
+    token_to_idx: Dict[str, int],
+) -> Dict[str, Any]:
+    vit_config = load_vit_config(model)
+    imgs: List[str] = []
+    vespa_sim_maps: List[str] = []
+    for single_result in result["root"]["children"]:
+        img = single_result["fields"]["full_image"]
+        if img:
+            imgs.append(img)
+        vespa_sim_map = single_result["fields"].get("summaryfeatures", None)
+        if vespa_sim_map:
+            vespa_sim_maps.append(vespa_sim_map)
+    sim_map_imgs = gen_similarity_maps(
+        model=model,
+        processor=processor,
+        device=model.device,
+        vit_config=vit_config,
+        query=query,
+        query_embs=q_embs,
+        token_idx_map=token_to_idx,
+        images=imgs,
+        vespa_sim_maps=vespa_sim_maps
     )
-    elt2["fields"]["page_number"] = 0
-    elt2["fields"]["image"] = "empty for now - is base64 encoded image"
-    result["root"]["children"].append(elt2)
+    for single_result, sim_map_dict in zip(result["root"]["children"], sim_map_imgs):
+        for token, sim_mapb64 in sim_map_dict.items():
+            single_result["fields"][f"sim_map_{token}"] = sim_mapb64
     return result
 
 
@@ -513,9 +530,24 @@ if __name__ == "__main__":
         / "assets"
         / "ConocoPhillips Sustainability Highlights - Nature (24-0976).png"
     )
-    gen_similarity_map(
-        model, processor, model.device, vit_config, query=query, image=image_filepath
+    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,
     )
-    result = get_result_dummy("dummy query")
-    print(result)
+    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/modelmanager.py

@@ -0,0 +1,22 @@
+from .colpali import load_model
+
+
+class ModelManager:
+    _instance = None
+    model = None
+    processor = None
+
+    @staticmethod
+    def get_instance():
+        if ModelManager._instance is None:
+            ModelManager._instance = ModelManager()
+            ModelManager._instance.initialize_model_and_processor()
+        return ModelManager._instance
+
+    def initialize_model_and_processor(self):
+        if self.model is None or self.processor is None:  # Ensure no reinitialization
+            self.model, self.processor = load_model()
+            if self.model is None or self.processor is None:
+                print("Failed to initialize model or processor at startup")
+            else:
+                print("Model and processor loaded at startup")

colpali-with-snippets/schemas/pdf_page.sd

@@ -0,0 +1,233 @@
+schema pdf_page {
+    document pdf_page {
+        field id type string {
+            indexing: summary | index
+            match {
+                word
+            }
+        }
+        field url type string {
+            indexing: summary | index
+        }
+        field title type string {
+            indexing: summary | index
+            index: enable-bm25
+            match {
+                text
+            }
+        }
+        field page_number type int {
+            indexing: summary | attribute
+        }
+        field image type raw {
+            indexing: summary
+        }
+        field full_image type raw {
+            indexing: summary
+        }
+        field text type string {
+            indexing: summary | index
+            index: enable-bm25
+            match {
+                text
+            }
+        }
+        field embedding type tensor<int8>(patch{}, v[16]) {
+            indexing: attribute | index
+            attribute {
+                distance-metric: hamming
+            }
+            index {
+                hnsw {
+                    max-links-per-node: 32
+                    neighbors-to-explore-at-insert: 400
+                }
+            }
+        }
+    }
+    fieldset default {
+        fields: title, url, page_number, text
+    }
+
+    document-summary default {
+        from-disk
+
+        summary text {
+            bolding: on
+        }
+
+        summary snippet {
+            source: text
+            dynamic
+        }
+    }
+
+    fieldset image {
+        fields: image
+    }
+
+    rank-profile bm25 {
+        first-phase {
+            expression: bm25(title) + bm25(text)
+        }
+    }
+
+    rank-profile default {
+        inputs {
+            query(qt) tensor<float>(querytoken{}, v[128])             
+        
+        }
+        function max_sim() {
+            expression {
+                
+                                    sum(
+                                        reduce(
+                                            sum(
+                                                query(qt) * unpack_bits(attribute(embedding)) , v
+                                            ),
+                                            max, patch
+                                        ),
+                                        querytoken
+                                    )
+                                
+            }
+        }
+        function similarities() {
+            expression  {
+                            sum(
+                                query(qt) * unpack_bits(attribute(embedding)), v
+                            )
+            }
+        }
+        function bm25_score() {
+            expression {
+                bm25(title) + bm25(text)
+            }
+        }
+        first-phase {
+            expression {
+                bm25_score
+            }
+        }
+        second-phase {
+            rerank-count: 10
+            expression {
+                max_sim
+            }
+        }
+        summary-features: similarities
+    }
+    rank-profile retrieval-and-rerank {
+        inputs {
+            query(rq0) tensor<int8>(v[16])             
+            query(rq1) tensor<int8>(v[16])             
+            query(rq2) tensor<int8>(v[16])             
+            query(rq3) tensor<int8>(v[16])             
+            query(rq4) tensor<int8>(v[16])             
+            query(rq5) tensor<int8>(v[16])             
+            query(rq6) tensor<int8>(v[16])             
+            query(rq7) tensor<int8>(v[16])             
+            query(rq8) tensor<int8>(v[16])             
+            query(rq9) tensor<int8>(v[16])             
+            query(rq10) tensor<int8>(v[16])             
+            query(rq11) tensor<int8>(v[16])             
+            query(rq12) tensor<int8>(v[16])             
+            query(rq13) tensor<int8>(v[16])             
+            query(rq14) tensor<int8>(v[16])             
+            query(rq15) tensor<int8>(v[16])             
+            query(rq16) tensor<int8>(v[16])             
+            query(rq17) tensor<int8>(v[16])             
+            query(rq18) tensor<int8>(v[16])             
+            query(rq19) tensor<int8>(v[16])             
+            query(rq20) tensor<int8>(v[16])             
+            query(rq21) tensor<int8>(v[16])             
+            query(rq22) tensor<int8>(v[16])             
+            query(rq23) tensor<int8>(v[16])             
+            query(rq24) tensor<int8>(v[16])             
+            query(rq25) tensor<int8>(v[16])             
+            query(rq26) tensor<int8>(v[16])             
+            query(rq27) tensor<int8>(v[16])             
+            query(rq28) tensor<int8>(v[16])             
+            query(rq29) tensor<int8>(v[16])             
+            query(rq30) tensor<int8>(v[16])             
+            query(rq31) tensor<int8>(v[16])             
+            query(rq32) tensor<int8>(v[16])             
+            query(rq33) tensor<int8>(v[16])             
+            query(rq34) tensor<int8>(v[16])             
+            query(rq35) tensor<int8>(v[16])             
+            query(rq36) tensor<int8>(v[16])             
+            query(rq37) tensor<int8>(v[16])             
+            query(rq38) tensor<int8>(v[16])             
+            query(rq39) tensor<int8>(v[16])             
+            query(rq40) tensor<int8>(v[16])             
+            query(rq41) tensor<int8>(v[16])             
+            query(rq42) tensor<int8>(v[16])             
+            query(rq43) tensor<int8>(v[16])             
+            query(rq44) tensor<int8>(v[16])             
+            query(rq45) tensor<int8>(v[16])             
+            query(rq46) tensor<int8>(v[16])             
+            query(rq47) tensor<int8>(v[16])             
+            query(rq48) tensor<int8>(v[16])             
+            query(rq49) tensor<int8>(v[16])             
+            query(rq50) tensor<int8>(v[16])             
+            query(rq51) tensor<int8>(v[16])             
+            query(rq52) tensor<int8>(v[16])             
+            query(rq53) tensor<int8>(v[16])             
+            query(rq54) tensor<int8>(v[16])             
+            query(rq55) tensor<int8>(v[16])             
+            query(rq56) tensor<int8>(v[16])             
+            query(rq57) tensor<int8>(v[16])             
+            query(rq58) tensor<int8>(v[16])             
+            query(rq59) tensor<int8>(v[16])             
+            query(rq60) tensor<int8>(v[16])             
+            query(rq61) tensor<int8>(v[16])             
+            query(rq62) tensor<int8>(v[16])             
+            query(rq63) tensor<int8>(v[16])             
+            query(qt) tensor<float>(querytoken{}, v[128])             
+            query(qtb) tensor<int8>(querytoken{}, v[16])             
+        
+        }
+        function max_sim() {
+            expression {
+                
+                                    sum(
+                                        reduce(
+                                            sum(
+                                                query(qt) * unpack_bits(attribute(embedding)) , v
+                                            ),
+                                            max, patch
+                                        ),
+                                        querytoken
+                                    )
+                                
+            }
+        }
+        function max_sim_binary() {
+            expression {
+                
+                                    sum(
+                                      reduce(
+                                        1/(1 + sum(
+                                            hamming(query(qtb), attribute(embedding)) ,v)
+                                        ),
+                                        max,
+                                        patch
+                                      ),
+                                      querytoken
+                                    )
+                                
+            }
+        }
+        first-phase {
+            expression {
+                max_sim_binary
+            }
+        }
+        second-phase {
+            rerank-count: 10
+            expression {
+                max_sim
+            }
+        }
+    }
+}

colpali-with-snippets/search/query-profiles/default.xml

@@ -0,0 +1,2 @@
+<query-profile id="default" type="root">
+</query-profile>

colpali-with-snippets/search/query-profiles/types/root.xml

@@ -0,0 +1,2 @@
+<query-profile-type id="root">
+</query-profile-type>

colpali-with-snippets/services.xml

@@ -0,0 +1,43 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<services version="1.0">
+    <container id="colpalidemo_container" version="1.0">
+        <search></search>
+        <document-api></document-api>
+        <document-processing></document-processing>
+        <clients>
+            <client id="mtls" permissions="read,write">
+                <certificate file="security/clients.pem" />
+            </client>
+            <client id="token_write" permissions="read,write">
+                <token id="colpalidemo_write" />
+            </client>
+            <client id="token_read" permissions="read">
+                <token id="colpalidemo_read" />
+            </client>
+        </clients>
+        <config name="container.qr-searchers">
+            <tag>
+                <bold>
+                    <open>&lt;strong&gt;</open>
+                    <close>&lt;/strong&gt;</close>
+                </bold>
+                <separator>...</separator>
+            </tag>
+        </config>
+    </container>
+    <content id="colpalidemo_content" version="1.0">
+        <redundancy>1</redundancy>
+        <documents>
+            <document type="pdf_page" mode="index"></document>
+        </documents>
+        <nodes>
+            <node distribution-key="0" hostalias="node1"></node>
+        </nodes>
+        <config name="vespa.config.search.summary.juniperrc">
+            <max_matches>2</max_matches>
+            <length>1000</length>
+            <surround_max>500</surround_max>
+            <min_length>300</min_length>
+        </config>
+    </content>
+</services>

colpalidemo/schemas/pdf_page.sd

@@ -22,6 +22,9 @@ schema pdf_page {
         field image type raw {
             indexing: summary
         }
+        field full_image type raw {
+            indexing: summary
+        }
         field text type string {
             indexing: summary | index
             index: enable-bm25
@@ -43,7 +46,10 @@ schema pdf_page {
         }
     }
     fieldset default {
-        fields: title, text
+        fields: title, url, page_number, text
+    }
+    fieldset image {
+        fields: image
     }
     rank-profile default {
         inputs {

deploy_vespa_app.py

@@ -16,6 +16,7 @@ from vespa.package import (
 )
 from vespa.deployment import VespaCloud
 import os
+from pathlib import Path
 
 
 def main():
@@ -60,6 +61,7 @@ def main():
                     name="page_number", type="int", indexing=["summary", "attribute"]
                 ),
                 Field(name="image", type="raw", indexing=["summary"]),
+                Field(name="full_image", type="raw", indexing=["summary"]),
                 Field(
                     name="text",
                     type="string",
@@ -190,10 +192,12 @@ def main():
         tenant=tenant_name,
         application=vespa_app_name,
         key_content=vespa_team_api_key,
-        application_package=vespa_application_package,
+        application_root="colpali-with-snippets",
+        #application_package=vespa_application_package,
     )
 
-    app = vespa_cloud.deploy()
+    #app = vespa_cloud.deploy()
+    vespa_cloud.deploy_from_disk("default", "colpali-with-snippets")
 
     # Output the endpoint URL
     endpoint_url = vespa_cloud.get_token_endpoint()

feed_vespa.py

@@ -159,6 +159,7 @@ def main():
                 base_64_image = get_base64_image(
                     scale_image(image, 640), add_url_prefix=False
                 )
+                base_64_full_image = get_base64_image(image, add_url_prefix=False)
                 embedding_dict = dict()
                 for idx, patch_embedding in enumerate(embedding):
                     binary_vector = (
@@ -178,6 +179,7 @@ def main():
                         "title": title,
                         "page_number": page_number,
                         "image": base_64_image,
+                        "full_image": base_64_full_image,
                         "text": page_text,
                         "embedding": embedding_dict,
                     },

frontend/app.py

@@ -1,26 +1,59 @@
 from urllib.parse import quote_plus
+from typing import Optional
 
-from fasthtml.components import Div, H1, P, Img, H2, Form, Span
-from fasthtml.xtend import Script, A
+from fasthtml.components import H1, H2, Div, Form, Img, P, Span, NotStr
+from fasthtml.xtend import A, Script
 from lucide_fasthtml import Lucide
-from shad4fast import Button, Input, Badge
+from shad4fast import Badge, Button, Input, Label, RadioGroup, RadioGroupItem
 
-
-def check_input_script():
-    return Script(
-        """
+# JavaScript to check the input value and enable/disable the search button and radio buttons
+check_input_script = Script(
+    """
         window.onload = function() {
             const input = document.getElementById('search-input');
             const button = document.querySelector('[data-button="search-button"]');
-            function checkInputValue() { button.disabled = input.value.trim() === ""; }
-            input.addEventListener('input', checkInputValue);
-            checkInputValue();
+            const radioGroupItems = document.querySelectorAll('button[data-ref="radio-item"]');  // Get all radio buttons
+            
+            function checkInputValue() {
+                const isInputEmpty = input.value.trim() === "";
+                button.disabled = isInputEmpty;  // Disable the submit button
+                radioGroupItems.forEach(item => {
+                    item.disabled = isInputEmpty;  // Disable/enable the radio buttons
+                });
+            }
+
+            input.addEventListener('input', checkInputValue);  // Listen for input changes
+            checkInputValue();  // Initial check when the page loads
         };
-        """
-    )
+    """
+)
+
+# JavaScript to handle the image swapping, reset button, and active class toggling
+image_swapping = Script(
+    """
+    document.addEventListener('click', function (e) {
+        if (e.target.classList.contains('sim-map-button') || e.target.classList.contains('reset-button')) {
+            const newSrc = e.target.getAttribute('data-image-src');
+            const img = e.target.closest('.relative').querySelector('.result-image');
+            img.src = newSrc;
+
+            // Remove 'active' class from previously active button
+            const activeButton = document.querySelector('.sim-map-button.active');
+            if (activeButton) {
+                activeButton.classList.remove('active');
+            }
 
+            // Add 'active' class to the clicked button (if it's a sim-map button)
+            if (e.target.classList.contains('sim-map-button')) {
+                e.target.classList.add('active');
+            }
+        }
+    });
+    """
+)
 
-def SearchBox(with_border=False, query_value=""):
+
+def SearchBox(with_border=False, query_value="", ranking_value="nn+colpali"):
     grid_cls = "grid gap-2 items-center p-3 bg-muted/80 dark:bg-muted/40 w-full"
 
     if with_border:
@@ -41,7 +74,30 @@ def SearchBox(with_border=False, query_value=""):
             cls="relative",
         ),
         Div(
-            Span("controls", cls="text-muted-foreground"),
+            Div(
+                Span("Ranking by:", cls="text-muted-foreground text-xs font-semibold"),
+                RadioGroup(
+                    Div(
+                        RadioGroupItem(value="nn+colpali", id="nn+colpali"),
+                        Label("nn+colpali", htmlFor="nn+colpali"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    Div(
+                        RadioGroupItem(value="bm25+colpali", id="bm25+colpali"),
+                        Label("bm25+colpali", htmlFor="bm25+colpali"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    Div(
+                        RadioGroupItem(value="bm25", id="bm25"),
+                        Label("bm25", htmlFor="bm25"),
+                        cls="flex items-center space-x-2",
+                    ),
+                    name="ranking",
+                    default_value=ranking_value,
+                    cls="grid-flow-col gap-x-5 text-muted-foreground",
+                ),
+                cls="grid grid-flow-col items-center gap-x-3 border border-input px-3 rounded-sm",
+            ),
             Button(
                 Lucide(icon="arrow-right", size="21"),
                 size="sm",
@@ -51,23 +107,23 @@ def SearchBox(with_border=False, query_value=""):
             ),
             cls="flex justify-between",
         ),
-        check_input_script(),
-        action=f"/search?query={quote_plus(query_value)}",  # This takes the user to /search with the loading message
+        check_input_script,
+        action=f"/search?query={quote_plus(query_value)}&ranking={quote_plus(ranking_value)}",
         method="GET",
-        hx_get=f"/fetch_results?query={quote_plus(query_value)}",  # This fetches the results asynchronously
-        hx_trigger="load",  # Trigger this after the page loads
-        hx_target="#search-results",  # Update the search results div dynamically
-        hx_swap="outerHTML",  # Replace the search results div entirely
-        hx_indicator="#loading-indicator",  # Show the loading indicator while fetching results
+        hx_get=f"/fetch_results?query={quote_plus(query_value)}&ranking={quote_plus(ranking_value)}",
+        hx_trigger="load",
+        hx_target="#search-results",
+        hx_swap="outerHTML",
+        hx_indicator="#loading-indicator",
         cls=grid_cls,
     )
 
 
 def SampleQueries():
     sample_queries = [
-        "What is the future of energy storage?",
-        "What is sustainable energy?",
-        "How to reduce carbon emissions?",
+        "Percentage of non-fresh water as source?",
+        "Policies related to nature risk?",
+        "How much of produced water is recycled?",
     ]
 
     query_badges = []
@@ -83,7 +139,7 @@ def SampleQueries():
                         cls="flex gap-2 items-center",
                     ),
                     variant="outline",
-                    cls="text-base font-normal text-muted-foreground",
+                    cls="text-base font-normal text-muted-foreground hover:border-black dark:hover:border-white",
                 ),
                 href=f"/search?query={quote_plus(query)}",
                 cls="no-underline",
@@ -96,7 +152,7 @@ def SampleQueries():
 def Hero():
     return Div(
         H1(
-            "Vespa.Ai + ColPali",
+            "Vespa.ai + ColPali",
             cls="text-5xl md:text-7xl font-bold tracking-wide md:tracking-wider bg-clip-text text-transparent bg-gradient-to-r from-black to-gray-700 dark:from-white dark:to-gray-300 animate-fade-in",
         ),
         P(
@@ -121,12 +177,14 @@ def Home():
 
 def Search(request, search_results=[]):
     query_value = request.query_params.get("query", "").strip()
+    ranking_value = request.query_params.get("ranking", "nn+colpali")
+    print(
+        f"Search: Fetching results for query: {query_value}, ranking: {ranking_value}"
+    )
 
     return Div(
         Div(
-            SearchBox(
-                query_value=query_value
-            ),  # Pass the query value to pre-fill the SearchBox
+            SearchBox(query_value=query_value, ranking_value=ranking_value),
             Div(
                 LoadingMessage(),  # Show the loading message initially
                 id="search-results",  # This will be replaced by the search results
@@ -145,7 +203,7 @@ def LoadingMessage():
     )
 
 
-def SearchResult(results=[]):
+def SearchResult(results: list, query_id: Optional[str] = None):
     if not results:
         return Div(
             P(
@@ -159,35 +217,95 @@ def SearchResult(results=[]):
     result_items = []
     for result in results:
         fields = result["fields"]  # Extract the 'fields' part of each result
-        base64_image = (
-            f"data:image/jpeg;base64,{fields['image']}"  # Format base64 image
+        full_image_base64 = f"data:image/jpeg;base64,{fields['full_image']}"
+
+        # Filter sim_map fields that are words with 4 or more characters
+        sim_map_fields = {
+            key: value
+            for key, value in fields.items()
+            if key.startswith("sim_map_") and len(key.split("_")[-1]) >= 4
+        }
+
+        # Generate buttons for the sim_map fields
+        sim_map_buttons = []
+        for key, value in sim_map_fields.items():
+            sim_map_base64 = f"data:image/jpeg;base64,{value}"
+            sim_map_buttons.append(
+                Button(
+                    key.split("_")[-1],
+                    size="sm",
+                    data_image_src=sim_map_base64,
+                    cls="sim-map-button pointer-events-auto font-mono text-xs h-5 rounded-none px-2",
+                )
+            )
+
+        # Add "Reset Image" button to restore the full image
+        reset_button = Button(
+            "Reset",
+            variant="outline",
+            size="sm",
+            data_image_src=full_image_base64,
+            cls="reset-button pointer-events-auto font-mono text-xs h-5 rounded-none px-2",
+        )
+        # Add "Tokens" button - this has no action, just a placeholder
+        tokens_button = Button(
+            Lucide(icon="images", size="15"),
+            "Tokens",
+            size="sm",
+            cls="tokens-button flex gap-[3px] font-bold pointer-events-none font-mono text-xs h-5 rounded-none px-2",
         )
-        # Print the fields that start with 'sim_map'
-        for key, value in fields.items():
-            if key.startswith("sim_map"):
-                print(f"{key}")
         result_items.append(
             Div(
                 Div(
-                    Img(src=base64_image, alt=fields["title"], cls="max-w-full h-auto"),
-                    cls="bg-background px-3 py-5",
+                    Div(
+                        tokens_button,
+                        *sim_map_buttons,
+                        reset_button,
+                        cls="flex flex-wrap gap-px w-full  pointer-events-none",
+                    ),
+                    Img(
+                        src=full_image_base64,
+                        alt=fields["title"],
+                        cls="result-image max-w-full h-auto",
+                    ),
+                    cls="relative grid gap-px content-start bg-background px-3 py-5",
                 ),
                 Div(
                     Div(
                         H2(fields["title"], cls="text-xl font-semibold"),
                         P(
-                            fields["text"], cls="text-muted-foreground"
-                        ),  # Use the URL as the description
+                            "Page " + str(fields["page_number"]),
+                            cls="text-muted-foreground",
+                        ),
+                        P(
+                            "Relevance score: " + str(result["relevance"]),
+                            cls="text-muted-foreground",
+                        ),
+                        P(NotStr(fields["snippet"]), cls="text-muted-foreground"),
+                        P(NotStr(fields["text"]), cls="text-muted-foreground"),
                         cls="text-sm grid gap-y-4",
                     ),
-                    cls="bg-background px-3 py-5",
+                    cls="bg-background px-3 py-5 hidden md:block",
                 ),
-                cls="grid grid-cols-subgrid col-span-2",
+                cls="grid grid-cols-1 md:grid-cols-2 col-span-2",
             )
         )
 
-    return Div(
-        *result_items,
-        cls="grid grid-cols-2 gap-px bg-border",
-        id="search-results",  # This will be the target for HTMX updates
-    )
+    if query_id is not None:
+        return Div(
+            *result_items,
+            image_swapping,
+            hx_get=f"/updated_search_results?query_id={query_id}",
+            hx_trigger="every 1s",
+            hx_target="#search-results",
+            hx_swap="outerHTML",
+            id="search-results",
+            cls="grid grid-cols-2 gap-px bg-border",
+        )
+    else:
+        return Div(
+            *result_items,
+            image_swapping,
+            id="search-results",
+            cls="grid grid-cols-2 gap-px bg-border",
+        )

globals.css

@@ -155,3 +155,17 @@
 .animate-slide-up {
     animation: slide-up 1s ease-out forwards;
 }
+
+.sim-map-button.active {
+    background-color: #61D790;
+    color: #2E2F27;
+
+    &:hover {
+        background-color: #61D790;
+    }
+}
+
+.tokens-button {
+    background-color: #B7E2F1;
+    color: #2E2F27;
+}

icons.py

@@ -1 +1 @@
-ICONS = {"chevrons-right": "<path d=\"m6 17 5-5-5-5\"></path><path d=\"m13 17 5-5-5-5\"></path>", "moon": "<path d=\"M12 3a6 6 0 0 0 9 9 9 9 0 1 1-9-9Z\"></path>", "sun": "<circle cx=\"12\" cy=\"12\" r=\"4\"></circle><path d=\"M12 2v2\"></path><path d=\"M12 20v2\"></path><path d=\"m4.93 4.93 1.41 1.41\"></path><path d=\"m17.66 17.66 1.41 1.41\"></path><path d=\"M2 12h2\"></path><path d=\"M20 12h2\"></path><path d=\"m6.34 17.66-1.41 1.41\"></path><path d=\"m19.07 4.93-1.41 1.41\"></path>", "github": "<path d=\"M15 22v-4a4.8 4.8 0 0 0-1-3.5c3 0 6-2 6-5.5.08-1.25-.27-2.48-1-3.5.28-1.15.28-2.35 0-3.5 0 0-1 0-3 1.5-2.64-.5-5.36-.5-8 0C6 2 5 2 5 2c-.3 1.15-.3 2.35 0 3.5A5.403 5.403 0 0 0 4 9c0 3.5 3 5.5 6 5.5-.39.49-.68 1.05-.85 1.65-.17.6-.22 1.23-.15 1.85v4\"></path><path d=\"M9 18c-4.51 2-5-2-7-2\"></path>", "slack": "<rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"13\" y=\"2\"></rect><path d=\"M19 8.5V10h1.5A1.5 1.5 0 1 0 19 8.5\"></path><rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"8\" y=\"14\"></rect><path d=\"M5 15.5V14H3.5A1.5 1.5 0 1 0 5 15.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"14\" y=\"13\"></rect><path d=\"M15.5 19H14v1.5a1.5 1.5 0 1 0 1.5-1.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"2\" y=\"8\"></rect><path d=\"M8.5 5H10V3.5A1.5 1.5 0 1 0 8.5 5\"></path>", "settings": "<path d=\"M12.22 2h-.44a2 2 0 0 0-2 2v.18a2 2 0 0 1-1 1.73l-.43.25a2 2 0 0 1-2 0l-.15-.08a2 2 0 0 0-2.73.73l-.22.38a2 2 0 0 0 .73 2.73l.15.1a2 2 0 0 1 1 1.72v.51a2 2 0 0 1-1 1.74l-.15.09a2 2 0 0 0-.73 2.73l.22.38a2 2 0 0 0 2.73.73l.15-.08a2 2 0 0 1 2 0l.43.25a2 2 0 0 1 1 1.73V20a2 2 0 0 0 2 2h.44a2 2 0 0 0 2-2v-.18a2 2 0 0 1 1-1.73l.43-.25a2 2 0 0 1 2 0l.15.08a2 2 0 0 0 2.73-.73l.22-.39a2 2 0 0 0-.73-2.73l-.15-.08a2 2 0 0 1-1-1.74v-.5a2 2 0 0 1 1-1.74l.15-.09a2 2 0 0 0 .73-2.73l-.22-.38a2 2 0 0 0-2.73-.73l-.15.08a2 2 0 0 1-2 0l-.43-.25a2 2 0 0 1-1-1.73V4a2 2 0 0 0-2-2z\"></path><circle cx=\"12\" cy=\"12\" r=\"3\"></circle>", "arrow-right": "<path d=\"M5 12h14\"></path><path d=\"m12 5 7 7-7 7\"></path>", "search": "<circle cx=\"11\" cy=\"11\" r=\"8\"></circle><path d=\"m21 21-4.3-4.3\"></path>", "file-search": "<path d=\"M14 2v4a2 2 0 0 0 2 2h4\"></path><path d=\"M4.268 21a2 2 0 0 0 1.727 1H18a2 2 0 0 0 2-2V7l-5-5H6a2 2 0 0 0-2 2v3\"></path><path d=\"m9 18-1.5-1.5\"></path><circle cx=\"5\" cy=\"14\" r=\"3\"></circle>", "message-circle-question": "<path d=\"M7.9 20A9 9 0 1 0 4 16.1L2 22Z\"></path><path d=\"M9.09 9a3 3 0 0 1 5.83 1c0 2-3 3-3 3\"></path><path d=\"M12 17h.01\"></path>", "text-search": "<path d=\"M21 6H3\"></path><path d=\"M10 12H3\"></path><path d=\"M10 18H3\"></path><circle cx=\"17\" cy=\"15\" r=\"3\"></circle><path d=\"m21 19-1.9-1.9\"></path>"}
+ICONS = {"chevrons-right": "<path d=\"m6 17 5-5-5-5\"></path><path d=\"m13 17 5-5-5-5\"></path>", "moon": "<path d=\"M12 3a6 6 0 0 0 9 9 9 9 0 1 1-9-9Z\"></path>", "sun": "<circle cx=\"12\" cy=\"12\" r=\"4\"></circle><path d=\"M12 2v2\"></path><path d=\"M12 20v2\"></path><path d=\"m4.93 4.93 1.41 1.41\"></path><path d=\"m17.66 17.66 1.41 1.41\"></path><path d=\"M2 12h2\"></path><path d=\"M20 12h2\"></path><path d=\"m6.34 17.66-1.41 1.41\"></path><path d=\"m19.07 4.93-1.41 1.41\"></path>", "github": "<path d=\"M15 22v-4a4.8 4.8 0 0 0-1-3.5c3 0 6-2 6-5.5.08-1.25-.27-2.48-1-3.5.28-1.15.28-2.35 0-3.5 0 0-1 0-3 1.5-2.64-.5-5.36-.5-8 0C6 2 5 2 5 2c-.3 1.15-.3 2.35 0 3.5A5.403 5.403 0 0 0 4 9c0 3.5 3 5.5 6 5.5-.39.49-.68 1.05-.85 1.65-.17.6-.22 1.23-.15 1.85v4\"></path><path d=\"M9 18c-4.51 2-5-2-7-2\"></path>", "slack": "<rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"13\" y=\"2\"></rect><path d=\"M19 8.5V10h1.5A1.5 1.5 0 1 0 19 8.5\"></path><rect height=\"8\" rx=\"1.5\" width=\"3\" x=\"8\" y=\"14\"></rect><path d=\"M5 15.5V14H3.5A1.5 1.5 0 1 0 5 15.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"14\" y=\"13\"></rect><path d=\"M15.5 19H14v1.5a1.5 1.5 0 1 0 1.5-1.5\"></path><rect height=\"3\" rx=\"1.5\" width=\"8\" x=\"2\" y=\"8\"></rect><path d=\"M8.5 5H10V3.5A1.5 1.5 0 1 0 8.5 5\"></path>", "settings": "<path d=\"M12.22 2h-.44a2 2 0 0 0-2 2v.18a2 2 0 0 1-1 1.73l-.43.25a2 2 0 0 1-2 0l-.15-.08a2 2 0 0 0-2.73.73l-.22.38a2 2 0 0 0 .73 2.73l.15.1a2 2 0 0 1 1 1.72v.51a2 2 0 0 1-1 1.74l-.15.09a2 2 0 0 0-.73 2.73l.22.38a2 2 0 0 0 2.73.73l.15-.08a2 2 0 0 1 2 0l.43.25a2 2 0 0 1 1 1.73V20a2 2 0 0 0 2 2h.44a2 2 0 0 0 2-2v-.18a2 2 0 0 1 1-1.73l.43-.25a2 2 0 0 1 2 0l.15.08a2 2 0 0 0 2.73-.73l.22-.39a2 2 0 0 0-.73-2.73l-.15-.08a2 2 0 0 1-1-1.74v-.5a2 2 0 0 1 1-1.74l.15-.09a2 2 0 0 0 .73-2.73l-.22-.38a2 2 0 0 0-2.73-.73l-.15.08a2 2 0 0 1-2 0l-.43-.25a2 2 0 0 1-1-1.73V4a2 2 0 0 0-2-2z\"></path><circle cx=\"12\" cy=\"12\" r=\"3\"></circle>", "arrow-right": "<path d=\"M5 12h14\"></path><path d=\"m12 5 7 7-7 7\"></path>", "search": "<circle cx=\"11\" cy=\"11\" r=\"8\"></circle><path d=\"m21 21-4.3-4.3\"></path>", "file-search": "<path d=\"M14 2v4a2 2 0 0 0 2 2h4\"></path><path d=\"M4.268 21a2 2 0 0 0 1.727 1H18a2 2 0 0 0 2-2V7l-5-5H6a2 2 0 0 0-2 2v3\"></path><path d=\"m9 18-1.5-1.5\"></path><circle cx=\"5\" cy=\"14\" r=\"3\"></circle>", "message-circle-question": "<path d=\"M7.9 20A9 9 0 1 0 4 16.1L2 22Z\"></path><path d=\"M9.09 9a3 3 0 0 1 5.83 1c0 2-3 3-3 3\"></path><path d=\"M12 17h.01\"></path>", "text-search": "<path d=\"M21 6H3\"></path><path d=\"M10 12H3\"></path><path d=\"M10 18H3\"></path><circle cx=\"17\" cy=\"15\" r=\"3\"></circle><path d=\"m21 19-1.9-1.9\"></path>", "maximize": "<path d=\"M8 3H5a2 2 0 0 0-2 2v3\"></path><path d=\"M21 8V5a2 2 0 0 0-2-2h-3\"></path><path d=\"M3 16v3a2 2 0 0 0 2 2h3\"></path><path d=\"M16 21h3a2 2 0 0 0 2-2v-3\"></path>", "expand": "<path d=\"m21 21-6-6m6 6v-4.8m0 4.8h-4.8\"></path><path d=\"M3 16.2V21m0 0h4.8M3 21l6-6\"></path><path d=\"M21 7.8V3m0 0h-4.8M21 3l-6 6\"></path><path d=\"M3 7.8V3m0 0h4.8M3 3l6 6\"></path>", "fullscreen": "<path d=\"M3 7V5a2 2 0 0 1 2-2h2\"></path><path d=\"M17 3h2a2 2 0 0 1 2 2v2\"></path><path d=\"M21 17v2a2 2 0 0 1-2 2h-2\"></path><path d=\"M7 21H5a2 2 0 0 1-2-2v-2\"></path><rect height=\"8\" rx=\"1\" width=\"10\" x=\"7\" y=\"8\"></rect>", "images": "<path d=\"M18 22H4a2 2 0 0 1-2-2V6\"></path><path d=\"m22 13-1.296-1.296a2.41 2.41 0 0 0-3.408 0L11 18\"></path><circle cx=\"12\" cy=\"8\" r=\"2\"></circle><rect height=\"16\" rx=\"2\" width=\"16\" x=\"6\" y=\"2\"></rect>", "circle": "<circle cx=\"12\" cy=\"12\" r=\"10\"></circle>"}

main.py

@@ -1,14 +1,23 @@
 import asyncio
-import json
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
 
 from fasthtml.common import *
 from shad4fast import *
 from vespa.application import Vespa
+import time
 
-from backend.colpali import load_model, get_result_from_query
+from backend.colpali import (
+    get_result_from_query,
+    get_query_embeddings_and_token_map,
+    add_sim_maps_to_result,
+)
 from backend.vespa_app import get_vespa_app
-from frontend.app import Home, Search, SearchResult, SearchBox
+from backend.cache import LRUCache
+from backend.modelmanager import ModelManager
+from frontend.app import Home, Search, SearchBox, SearchResult
 from frontend.layout import Layout
+import hashlib
 
 highlight_js_theme_link = Link(id="highlight-theme", rel="stylesheet", href="")
 highlight_js_theme = Script(src="/static/js/highlightjs-theme.js")
@@ -30,26 +39,12 @@ app, rt = fast_app(
 )
 vespa_app: Vespa = get_vespa_app()
 
+result_cache = LRUCache(max_size=20)  # Each result can be ~10MB
+thread_pool = ThreadPoolExecutor()
 
-class ModelManager:
-    _instance = None
-    model = None
-    processor = None
-
-    @staticmethod
-    def get_instance():
-        if ModelManager._instance is None:
-            ModelManager._instance = ModelManager()
-            ModelManager._instance.initialize_model_and_processor()
-        return ModelManager._instance
 
-    def initialize_model_and_processor(self):
-        if self.model is None or self.processor is None:  # Ensure no reinitialization
-            self.model, self.processor = load_model()
-            if self.model is None or self.processor is None:
-                print("Failed to initialize model or processor at startup")
-            else:
-                print("Model and processor loaded at startup")
+def generate_query_id(query):
+    return hashlib.md5(query.encode("utf-8")).hexdigest()
 
 
 @rt("/static/{filepath:path}")
@@ -64,15 +59,17 @@ def get():
 
 @rt("/search")
 def get(request):
-    # Extract the 'query' parameter from the URL using query_params
+    # Extract the 'query' and 'ranking' parameters from the URL
     query_value = request.query_params.get("query", "").strip()
+    ranking_value = request.query_params.get("ranking", "nn+colpali")
+    print("/search: Fetching results for ranking_value:", ranking_value)
 
     # Always render the SearchBox first
     if not query_value:
         # Show SearchBox and a message for missing query
         return Layout(
             Div(
-                SearchBox(query_value=query_value),
+                SearchBox(query_value=query_value, ranking_value=ranking_value),
                 Div(
                     P(
                         "No query provided. Please enter a query.",
@@ -89,39 +86,80 @@ def get(request):
 
 
 @rt("/fetch_results")
-def get(request, query: str, nn: bool = True):
-    # Check if the request came from HTMX; if not, redirect to /search
+async def get(request, query: str, nn: bool = True):
     if "hx-request" not in request.headers:
         return RedirectResponse("/search")
 
-    # Extract the 'query' parameter from the URL
+    # Extract ranking option from the request
+    ranking_value = request.query_params.get("ranking")
+    print(
+        f"/fetch_results: Fetching results for query: {query}, ranking: {ranking_value}"
+    )
+    # Generate a unique query_id based on the query and ranking value
+    query_id = generate_query_id(query + ranking_value)
 
     # Fetch model and processor
     manager = ModelManager.get_instance()
     model = manager.model
     processor = manager.processor
+    q_embs, token_to_idx = get_query_embeddings_and_token_map(processor, model, query)
 
+    start = time.perf_counter()
     # Fetch real search results from Vespa
-    result = asyncio.run(
-        get_result_from_query(
-            vespa_app,
-            processor=processor,
-            model=model,
-            query=query,
-            nn=nn,
-            gen_sim_map=False,
+    result = await get_result_from_query(
+        app=vespa_app,
+        processor=processor,
+        model=model,
+        query=query,
+        q_embs=q_embs,
+        token_to_idx=token_to_idx,
+        ranking=ranking_value,
+    )
+    end = time.perf_counter()
+    print(f"Search results fetched in {end - start:.2f} seconds, Vespa says searchtime was {result['timing']['searchtime']} seconds")
+    # Start generating the similarity map in the background
+    asyncio.create_task(
+        generate_similarity_map(
+            model, processor, query, q_embs, token_to_idx, result, query_id
         )
     )
-
-    # Extract search results from the result payload
     search_results = (
         result["root"]["children"]
         if "root" in result and "children" in result["root"]
         else []
     )
+    return SearchResult(search_results, query_id)
+
+
+async def generate_similarity_map(
+    model, processor, query, q_embs, token_to_idx, result, query_id
+):
+    loop = asyncio.get_event_loop()
+    sim_map_task = partial(
+        add_sim_maps_to_result,
+        result=result,
+        model=model,
+        processor=processor,
+        query=query,
+        q_embs=q_embs,
+        token_to_idx=token_to_idx,
+    )
+    sim_map_result = await loop.run_in_executor(thread_pool, sim_map_task)
+    result_cache.set(query_id, sim_map_result)
+
 
-    # Directly return the search results without the full page layout
-    return SearchResult(search_results)
+@app.get("/updated_search_results")
+async def updated_search_results(query_id: str):
+    data = result_cache.get(query_id)
+    if data is None:
+        return HTMLResponse(status_code=204)
+    search_results = (
+        data["root"]["children"]
+        if "root" in data and "children" in data["root"]
+        else []
+    )
+    updated_content = SearchResult(results=search_results, query_id=None)
+    return updated_content
 
 
 @rt("/app")
@@ -129,24 +167,6 @@ def get():
     return Layout(Div(P(f"Connected to Vespa at {vespa_app.url}"), cls="p-4"))
 
 
-@rt("/run_query")
-def get(query: str, nn: bool = False):
-    # dummy-function to avoid running the query every time
-    # result = get_result_dummy(query, nn)
-    # If we want to run real, uncomment the following lines
-    model, processor = get_model_and_processor()
-    result = asyncio.run(
-        get_result_from_query(
-            vespa_app, processor=processor, model=model, query=query, nn=nn
-        )
-    )
-    # model, processor = get_model_and_processor()
-    # result = asyncio.run(
-    #     get_result_from_query(vespa_app, processor=processor, model=model, query=query, nn=nn)
-    # )
-    return Layout(Div(H1("Result"), Pre(Code(json.dumps(result, indent=2))), cls="p-4"))
-
-
 if __name__ == "__main__":
     # ModelManager.get_instance()  # Initialize once at startup
     serve(port=7860)

output.css

@@ -1073,12 +1073,16 @@ body {
   resize: both;
 }
 
-.grid-cols-2 {
-  grid-template-columns: repeat(2, minmax(0, 1fr));
+.grid-flow-col {
+  grid-auto-flow: column;
 }
 
-.grid-cols-subgrid {
-  grid-template-columns: subgrid;
+.grid-cols-1 {
+  grid-template-columns: repeat(1, minmax(0, 1fr));
+}
+
+.grid-cols-2 {
+  grid-template-columns: repeat(2, minmax(0, 1fr));
 }
 
 .flex-col {
@@ -1089,6 +1093,14 @@ body {
   flex-direction: column-reverse;
 }
 
+.flex-wrap {
+  flex-wrap: wrap;
+}
+
+.content-start {
+  align-content: flex-start;
+}
+
 .items-center {
   align-items: center;
 }
@@ -1121,18 +1133,34 @@ body {
   gap: 2rem;
 }
 
+.gap-\[3px\] {
+  gap: 3px;
+}
+
 .gap-px {
   gap: 1px;
 }
 
-.gap-3 {
-  gap: 0.75rem;
+.gap-x-3 {
+  -moz-column-gap: 0.75rem;
+       column-gap: 0.75rem;
+}
+
+.gap-x-5 {
+  -moz-column-gap: 1.25rem;
+       column-gap: 1.25rem;
 }
 
 .gap-y-4 {
   row-gap: 1rem;
 }
 
+.space-x-2 > :not([hidden]) ~ :not([hidden]) {
+  --tw-space-x-reverse: 0;
+  margin-right: calc(0.5rem * var(--tw-space-x-reverse));
+  margin-left: calc(0.5rem * calc(1 - var(--tw-space-x-reverse)));
+}
+
 .space-x-3 > :not([hidden]) ~ :not([hidden]) {
   --tw-space-x-reverse: 0;
   margin-right: calc(0.75rem * var(--tw-space-x-reverse));
@@ -1193,6 +1221,10 @@ body {
   border-radius: calc(var(--radius) - 2px);
 }
 
+.rounded-none {
+  border-radius: 0px;
+}
+
 .rounded-sm {
   border-radius: calc(var(--radius) - 4px);
 }
@@ -1339,8 +1371,8 @@ body {
   padding: 0.25rem;
 }
 
-.p-16 {
-  padding: 4rem;
+.p-10 {
+  padding: 2.5rem;
 }
 
 .p-3 {
@@ -1359,8 +1391,9 @@ body {
   padding: 1px;
 }
 
-.p-10 {
-  padding: 2.5rem;
+.px-2 {
+  padding-left: 0.5rem;
+  padding-right: 0.5rem;
 }
 
 .px-2\.5 {
@@ -1448,6 +1481,10 @@ body {
   vertical-align: middle;
 }
 
+.font-mono {
+  font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
+}
+
 .text-2xl {
   font-size: 1.5rem;
   line-height: 2rem;
@@ -1496,14 +1533,14 @@ body {
   font-weight: 500;
 }
 
-.font-semibold {
-  font-weight: 600;
-}
-
 .font-normal {
   font-weight: 400;
 }
 
+.font-semibold {
+  font-weight: 600;
+}
+
 .leading-none {
   line-height: 1;
 }
@@ -1574,10 +1611,6 @@ body {
   color: transparent;
 }
 
-.underline {
-  text-decoration-line: underline;
-}
-
 .no-underline {
   text-decoration-line: none;
 }
@@ -1908,6 +1941,19 @@ body {
   animation: slide-up 1s ease-out forwards;
 }
 
+.sim-map-button.active {
+  background-color: #61D790;
+  color: #2E2F27;
+  &:hover {
+    background-color: #61D790;
+  }
+}
+
+.tokens-button {
+  background-color: #B7E2F1;
+  color: #2E2F27;
+}
+
 :root:has(.data-\[state\=open\]\:no-bg-scroll[data-state="open"]) {
   overflow: hidden;
 }
@@ -1964,13 +2010,9 @@ body {
   --tw-ring-offset-width: 2px;
 }
 
-.hover\:border-white:hover {
+.hover\:border-black:hover {
   --tw-border-opacity: 1;
-  border-color: rgb(255 255 255 / var(--tw-border-opacity));
-}
-
-.hover\:border-\[text-muted-foreground\]:hover {
-  border-color: text-muted-foreground;
+  border-color: rgb(0 0 0 / var(--tw-border-opacity));
 }
 
 .hover\:bg-accent:hover {
@@ -2001,10 +2043,6 @@ body {
   background-color: hsl(var(--secondary) / 0.8);
 }
 
-.hover\:bg-secondary:hover {
-  background-color: hsl(var(--secondary));
-}
-
 .hover\:text-accent-foreground:hover {
   color: hsl(var(--accent-foreground));
 }
@@ -2013,14 +2051,6 @@ body {
   color: hsl(var(--foreground));
 }
 
-.hover\:text-primary-foreground:hover {
-  color: hsl(var(--primary-foreground));
-}
-
-.hover\:text-muted-foreground:hover {
-  color: hsl(var(--muted-foreground));
-}
-
 .hover\:underline:hover {
   text-decoration-line: underline;
 }
@@ -2407,10 +2437,18 @@ body {
 }
 
 @media (min-width: 768px) {
+  .md\:block {
+    display: block;
+  }
+
   .md\:max-w-\[420px\] {
     max-width: 420px;
   }
 
+  .md\:grid-cols-2 {
+    grid-template-columns: repeat(2, minmax(0, 1fr));
+  }
+
   .md\:text-2xl {
     font-size: 1.5rem;
     line-height: 2rem;
@@ -2460,14 +2498,9 @@ body {
   --tw-gradient-to: #d1d5db var(--tw-gradient-to-position);
 }
 
-.dark\:hover\:border-black:hover:where(.dark, .dark *) {
-  --tw-border-opacity: 1;
-  border-color: rgb(0 0 0 / var(--tw-border-opacity));
-}
-
-.hover\:dark\:border-black:where(.dark, .dark *):hover {
+.dark\:hover\:border-white:hover:where(.dark, .dark *) {
   --tw-border-opacity: 1;
-  border-color: rgb(0 0 0 / var(--tw-border-opacity));
+  border-color: rgb(255 255 255 / var(--tw-border-opacity));
 }
 
 .\[\&\:has\(\[role\=checkbox\]\)\]\:pr-0:has([role=checkbox]) {

ruff.toml

@@ -0,0 +1,77 @@
+# Exclude a variety of commonly ignored directories.
+exclude = [
+    ".bzr",
+    ".direnv",
+    ".eggs",
+    ".git",
+    ".git-rewrite",
+    ".hg",
+    ".ipynb_checkpoints",
+    ".mypy_cache",
+    ".nox",
+    ".pants.d",
+    ".pyenv",
+    ".pytest_cache",
+    ".pytype",
+    ".ruff_cache",
+    ".svn",
+    ".tox",
+    ".venv",
+    ".vscode",
+    "__pypackages__",
+    "_build",
+    "buck-out",
+    "build",
+    "dist",
+    "node_modules",
+    "site-packages",
+    "venv",
+]
+
+# Same as Black.
+line-length = 88
+indent-width = 4
+
+# Assume Python 3.8
+target-version = "py38"
+
+[lint]
+# Enable Pyflakes (`F`) and a subset of the pycodestyle (`E`)  codes by default.
+# Unlike Flake8, Ruff doesn't enable pycodestyle warnings (`W`) or
+# McCabe complexity (`C901`) by default.
+select = ["E4", "E7", "E9", "F"]
+ignore = []
+
+# Allow fix for all enabled rules (when `--fix`) is provided.
+fixable = ["ALL"]
+unfixable = []
+
+# Allow unused variables when underscore-prefixed.
+dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
+
+[format]
+# Like Black, use double quotes for strings.
+quote-style = "double"
+
+# Like Black, indent with spaces, rather than tabs.
+indent-style = "space"
+
+# Like Black, respect magic trailing commas.
+skip-magic-trailing-comma = false
+
+# Like Black, automatically detect the appropriate line ending.
+line-ending = "auto"
+
+# Enable auto-formatting of code examples in docstrings. Markdown,
+# reStructuredText code/literal blocks and doctests are all supported.
+#
+# This is currently disabled by default, but it is planned for this
+# to be opt-out in the future.
+docstring-code-format = false
+
+# Set the line length limit used when formatting code snippets in
+# docstrings.
+#
+# This only has an effect when the `docstring-code-format` setting is
+# enabled.
+docstring-code-line-length = "dynamic"