add mapping between metadata (qdrant collections)

src/colpali/visual_search.py

@@ -8,10 +8,10 @@ All dependencies are now within src/colpali/ - no external colpali_colab_package
 """
 
 import logging
+import os
 from typing import List, Dict, Any, Optional
-
-import torch
 import numpy as np
+import torch
 from qdrant_client import QdrantClient
 
 # Import from local src/colpali modules (no external dependencies)
@@ -21,8 +21,10 @@ from src.colpali.search import VisualDocumentSearch
 # Import device detection utility
 from src.utils import get_device_for_colpali
 
+# Import filename mapping for v1 -> visual collection translation
+from src.config.visual_filename_mapping import v1_filenames_to_visual
+
 logger = logging.getLogger(__name__)
-DEFAULT_MODEL = "colSmol-500M"
 
 
 class VisualSearchResult:
@@ -55,8 +57,8 @@ class VisualSearchAdapter:
         self,
         qdrant_url: str,
         qdrant_api_key: str,
-        collection_name: str = DEFAULT_MODEL,
-        model_name: str = f"vidore/{DEFAULT_MODEL}",
+        collection_name: str = None,  # Will use QDRANT_COLLECTION_VISUAL env var or default
+        model_name: str = "vidore/colSmol-500M",
         device: str = None,
         batch_size: int = 4
     ):
@@ -66,13 +68,17 @@ class VisualSearchAdapter:
         Args:
             qdrant_url: Qdrant cluster URL
             qdrant_api_key: Qdrant API key
-            collection_name: Name of the collection with visual embeddings
+            collection_name: Name of the collection with visual embeddings (default from QDRANT_COLLECTION_VISUAL env var)
             model_name: ColPali model name
             device: Device to use (cuda/cpu/mps, auto-detected if None)
             batch_size: Batch size for embedding generation
         """
         logger.info("🎨 Initializing Visual Search Adapter...")
         
+        # Get collection name from env var if not provided
+        if collection_name is None:
+            collection_name = os.environ.get("QDRANT_COLLECTION_VISUAL", "colSmol-500M-v2")
+        
         # Auto-detect device using utility function
         if device is None:
             device = get_device_for_colpali()
@@ -134,11 +140,15 @@ class VisualSearchAdapter:
         """
         logger.info(f"🔍 Visual search: '{query}' (top_k={top_k}, strategy={search_strategy})")
         
-        # Generate query embedding
-        query_embedding = self.processor.embed_query(query)
+        # Generate query embedding (filter special tokens by default)
+        query_embedding = self.processor.embed_query(query, filter_special_tokens=True)
         
         # Store for saliency generation
         self.last_query_embedding = query_embedding
+        self.last_query_text = query  # Store query text for word selection
+        # Store processed query info for accurate word-to-token mapping
+        self.last_input_ids = getattr(self.processor, 'last_input_ids', None)
+        self.last_attention_mask = getattr(self.processor, 'last_attention_mask', None)
         
         # Convert filters to Qdrant format
         filter_params = {}
@@ -154,7 +164,14 @@ class VisualSearchAdapter:
             if 'districts' in filters and filters['districts']:
                 filter_params['district'] = filters['districts']
             if 'filenames' in filters and filters['filenames']:
-                filter_params['filename'] = filters['filenames']
+                v1_filenames = filters['filenames']
+                visual_filenames = v1_filenames_to_visual(v1_filenames)
+                if visual_filenames:
+                    logger.info(f"🔄 Filename translation: {v1_filenames} -> {visual_filenames}")
+                    filter_params['filename'] = visual_filenames
+                else:
+                    logger.warning(f"⚠️ No visual filename mappings found for: {v1_filenames}")
+                    filter_params['filename'] = v1_filenames
             if 'has_text' in filters:
                 filter_params['has_text'] = filters['has_text']
         
@@ -209,7 +226,7 @@ class VisualSearchAdapter:
 def create_visual_search_adapter(
     qdrant_url: Optional[str] = None,
     qdrant_api_key: Optional[str] = None,
-    collection_name: str = DEFAULT_MODEL
+    collection_name: Optional[str] = None
 ) -> VisualSearchAdapter:
     """
     Factory function to create a visual search adapter.
@@ -217,7 +234,7 @@ def create_visual_search_adapter(
     Args:
         qdrant_url: Qdrant URL (reads from env if not provided)
         qdrant_api_key: Qdrant API key (reads from env if not provided)
-        collection_name: Collection name
+        collection_name: Collection name (reads from QDRANT_COLLECTION_VISUAL env var if not provided)
     
     Returns:
         Initialized VisualSearchAdapter
@@ -228,6 +245,8 @@ def create_visual_search_adapter(
         qdrant_url = os.environ.get("QDRANT_URL")
     if qdrant_api_key is None:
         qdrant_api_key = os.environ.get("QDRANT_API_KEY")
+    if collection_name is None:
+        collection_name = os.environ.get("QDRANT_COLLECTION_VISUAL", "colSmol-500M-v2")
     
     if not qdrant_url or not qdrant_api_key:
         raise ValueError("QDRANT_URL and QDRANT_API_KEY must be provided or set in environment")