add cascaseded retrievers

visual_rag/__init__.py

@@ -61,6 +61,16 @@ try:
 except ImportError:
     TwoStageRetriever = None
 
+try:
+    from visual_rag.retrieval.multi_vector import MultiVectorRetriever
+except ImportError:
+    MultiVectorRetriever = None
+
+try:
+    from visual_rag.qdrant_admin import QdrantAdmin
+except ImportError:
+    QdrantAdmin = None
+
 # Config utilities (always available)
 from visual_rag.config import load_config, get, get_section
 
@@ -74,6 +84,8 @@ __all__ = [
     "QdrantIndexer",
     "CloudinaryUploader",
     "TwoStageRetriever",
+    "MultiVectorRetriever",
+    "QdrantAdmin",
     
     # Config utilities
     "load_config",

visual_rag/cli/main.py

@@ -24,6 +24,7 @@ import argparse
 import logging
 from pathlib import Path
 from typing import Optional
+from urllib.parse import urlparse
 
 from dotenv import load_dotenv
 
@@ -94,13 +95,31 @@ def cmd_process(args):
     model_name = args.model or config.get("model", {}).get("name", "vidore/colSmol-500M")
     collection_name = args.collection or config.get("qdrant", {}).get("collection_name", "visual_documents")
     
-    # Initialize embedder
+    torch_dtype = None
+    if args.torch_dtype != "auto":
+        import torch
+        torch_dtype = {
+            "float32": torch.float32,
+            "float16": torch.float16,
+            "bfloat16": torch.bfloat16,
+        }[args.torch_dtype]
+
     logger.info(f"🤖 Initializing embedder: {model_name}")
-    embedder = VisualEmbedder(model_name=model_name, batch_size=args.batch_size)
+    embedder = VisualEmbedder(
+        model_name=model_name,
+        batch_size=args.batch_size,
+        torch_dtype=torch_dtype,
+        processor_speed=str(getattr(args, "processor_speed", "fast")),
+    )
     
     # Initialize Qdrant indexer
-    qdrant_url = os.getenv("QDRANT_URL")
-    qdrant_api_key = os.getenv("QDRANT_API_KEY")
+    qdrant_url = os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
     
     if not qdrant_url:
         logger.error("❌ QDRANT_URL environment variable not set")
@@ -111,11 +130,42 @@ def cmd_process(args):
         url=qdrant_url,
         api_key=qdrant_api_key,
         collection_name=collection_name,
+        prefer_grpc=args.prefer_grpc,
+        vector_datatype=args.qdrant_vector_dtype,
     )
     
     # Create collection if needed
     indexer.create_collection(force_recreate=args.force_recreate)
-    indexer.create_payload_indexes()
+    inferred_fields = []
+    inferred_fields.append({"field": "filename", "type": "keyword"})
+    inferred_fields.append({"field": "page_number", "type": "integer"})
+    inferred_fields.append({"field": "has_text", "type": "bool"})
+
+    if metadata_mapping:
+        keys = set()
+        for _, meta in metadata_mapping.items():
+            if isinstance(meta, dict):
+                keys.update(meta.keys())
+        for k in sorted(keys):
+            if k in ("filename", "page_number", "has_text"):
+                continue
+            inferred_type = "keyword"
+            for _, meta in metadata_mapping.items():
+                if not isinstance(meta, dict):
+                    continue
+                v = meta.get(k)
+                if isinstance(v, bool):
+                    inferred_type = "bool"
+                    break
+                if isinstance(v, int):
+                    inferred_type = "integer"
+                    break
+                if isinstance(v, float):
+                    inferred_type = "float"
+                    break
+            inferred_fields.append({"field": k, "type": inferred_type})
+
+    indexer.create_payload_indexes(fields=inferred_fields)
     
     # Initialize Cloudinary uploader (optional)
     cloudinary_uploader = None
@@ -135,6 +185,9 @@ def cmd_process(args):
         metadata_mapping=metadata_mapping,
         config=config,
         embedding_strategy=args.strategy,
+        crop_empty=bool(getattr(args, "crop_empty", False)),
+        crop_empty_percentage_to_remove=float(getattr(args, "crop_empty_percentage_to_remove", 0.9)),
+        crop_empty_remove_page_number=bool(getattr(args, "crop_empty_remove_page_number", False)),
     )
     
     # Process PDFs
@@ -177,13 +230,18 @@ def cmd_process(args):
 def cmd_search(args):
     """Search documents."""
     from visual_rag import VisualEmbedder
-    from visual_rag.retrieval import TwoStageRetriever
+    from visual_rag.retrieval import TwoStageRetriever, SingleStageRetriever
     from qdrant_client import QdrantClient
     
     load_dotenv()
     
-    qdrant_url = os.getenv("QDRANT_URL")
-    qdrant_api_key = os.getenv("QDRANT_API_KEY")
+    qdrant_url = os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
     
     if not qdrant_url:
         logger.error("❌ QDRANT_URL not set")
@@ -191,11 +249,19 @@ def cmd_search(args):
     
     # Initialize
     logger.info(f"🤖 Loading model: {args.model}")
-    embedder = VisualEmbedder(model_name=args.model)
-    
+    embedder = VisualEmbedder(model_name=args.model, processor_speed=str(getattr(args, "processor_speed", "fast")))
+
     logger.info(f"🔌 Connecting to Qdrant")
-    client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
-    retriever = TwoStageRetriever(client, args.collection)
+    grpc_port = 6334 if args.prefer_grpc and urlparse(qdrant_url).port == 6333 else None
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        grpc_port=grpc_port,
+        check_compatibility=False,
+    )
+    two_stage = TwoStageRetriever(client, args.collection)
+    single_stage = SingleStageRetriever(client, args.collection)
     
     # Embed query
     logger.info(f"🔍 Query: {args.query}")
@@ -204,19 +270,43 @@ def cmd_search(args):
     # Build filter
     filter_obj = None
     if args.year or args.source or args.district:
-        filter_obj = retriever.build_filter(
+        filter_obj = two_stage.build_filter(
             year=args.year,
             source=args.source,
             district=args.district,
         )
     
     # Search
-    results = retriever.search(
-        query_embedding=query_embedding.numpy(),
-        top_k=args.top_k,
-        prefetch_k=args.prefetch_k,
-        filter_obj=filter_obj,
-    )
+    query_np = query_embedding.detach().cpu().numpy()
+    if args.strategy == "single_full":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="multi_vector",
+            filter_obj=filter_obj,
+        )
+    elif args.strategy == "single_tiles":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="tiles_maxsim",
+            filter_obj=filter_obj,
+        )
+    elif args.strategy == "single_global":
+        results = single_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            strategy="pooled_global",
+            filter_obj=filter_obj,
+        )
+    else:
+        results = two_stage.search(
+            query_embedding=query_np,
+            top_k=args.top_k,
+            prefetch_k=args.prefetch_k,
+            filter_obj=filter_obj,
+            stage1_mode=args.stage1_mode,
+        )
     
     # Display results
     logger.info(f"\n📊 Results ({len(results)}):")
@@ -245,14 +335,26 @@ def cmd_info(args):
     
     load_dotenv()
     
-    qdrant_url = os.getenv("QDRANT_URL")
-    qdrant_api_key = os.getenv("QDRANT_API_KEY")
+    qdrant_url = os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    qdrant_api_key = (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
     
     if not qdrant_url:
         logger.error("❌ QDRANT_URL not set")
         sys.exit(1)
     
-    client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
+    grpc_port = 6334 if args.prefer_grpc and urlparse(qdrant_url).port == 6333 else None
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        grpc_port=grpc_port,
+        check_compatibility=False,
+    )
     
     try:
         info = client.get_collection(args.collection)
@@ -345,6 +447,22 @@ Examples:
         "--no-cloudinary", action="store_true",
         help="Skip Cloudinary uploads"
     )
+    process_parser.add_argument(
+        "--crop-empty",
+        action="store_true",
+        help="Crop empty whitespace from page images before embedding (default: off).",
+    )
+    process_parser.add_argument(
+        "--crop-empty-percentage-to-remove",
+        type=float,
+        default=0.9,
+        help="Kept for traceability; currently does not affect cropping behavior (default: 0.9).",
+    )
+    process_parser.add_argument(
+        "--crop-empty-remove-page-number",
+        action="store_true",
+        help="If set, attempts to crop away the bottom region that contains sparse page numbers (default: off).",
+    )
     process_parser.add_argument(
         "--no-skip-existing", action="store_true",
         help="Process all pages even if they exist in Qdrant"
@@ -363,6 +481,41 @@ Examples:
         help="Embedding strategy: 'pooling' (NOVEL), 'standard' (BASELINE), "
              "'all' (embed once, store BOTH for comparison)"
     )
+    process_parser.add_argument(
+        "--torch-dtype",
+        type=str,
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+        help="Torch dtype for model weights (default: auto; CUDA->bfloat16, else float32).",
+    )
+    process_parser.add_argument(
+        "--qdrant-vector-dtype",
+        type=str,
+        default="float16",
+        choices=["float16", "float32"],
+        help="Datatype for vectors stored in Qdrant (default: float16).",
+    )
+    process_parser.add_argument(
+        "--processor-speed",
+        type=str,
+        default="fast",
+        choices=["fast", "slow", "auto"],
+        help="Processor implementation: fast (default, with fallback to slow), slow, or auto.",
+    )
+    process_grpc_group = process_parser.add_mutually_exclusive_group()
+    process_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    process_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
     process_parser.set_defaults(func=cmd_process)
     
     # =========================================================================
@@ -384,14 +537,31 @@ Examples:
         "--model", type=str, default="vidore/colSmol-500M",
         help="Model name"
     )
+    search_parser.add_argument(
+        "--processor-speed",
+        type=str,
+        default="fast",
+        choices=["fast", "slow", "auto"],
+        help="Processor implementation: fast (default, with fallback to slow), slow, or auto.",
+    )
     search_parser.add_argument(
         "--top-k", type=int, default=10,
         help="Number of results"
     )
+    search_parser.add_argument(
+        "--strategy", type=str, default="single_full",
+        choices=["single_full", "single_tiles", "single_global", "two_stage"],
+        help="Search strategy"
+    )
     search_parser.add_argument(
         "--prefetch-k", type=int, default=200,
         help="Prefetch candidates for two-stage retrieval"
     )
+    search_parser.add_argument(
+        "--stage1-mode", type=str, default="pooled_query_vs_tiles",
+        choices=["pooled_query_vs_tiles", "tokens_vs_tiles", "pooled_query_vs_global"],
+        help="Stage 1 mode for two-stage retrieval"
+    )
     search_parser.add_argument(
         "--year", type=int,
         help="Filter by year"
@@ -408,6 +578,20 @@ Examples:
         "--show-text", action="store_true",
         help="Show text snippets in results"
     )
+    search_grpc_group = search_parser.add_mutually_exclusive_group()
+    search_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    search_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
     search_parser.set_defaults(func=cmd_search)
     
     # =========================================================================
@@ -421,6 +605,20 @@ Examples:
         "--collection", type=str, default="visual_documents",
         help="Qdrant collection name"
     )
+    info_grpc_group = info_parser.add_mutually_exclusive_group()
+    info_grpc_group.add_argument(
+        "--prefer-grpc",
+        dest="prefer_grpc",
+        action="store_true",
+        default=True,
+        help="Use gRPC for Qdrant client (recommended).",
+    )
+    info_grpc_group.add_argument(
+        "--no-prefer-grpc",
+        dest="prefer_grpc",
+        action="store_false",
+        help="Disable gRPC for Qdrant client.",
+    )
     info_parser.set_defaults(func=cmd_info)
     
     # Parse and execute

visual_rag/config.py

@@ -18,6 +18,19 @@ logger = logging.getLogger(__name__)
 _config_cache: Optional[Dict[str, Any]] = None
 
 
+def _env_qdrant_url() -> Optional[str]:
+    return os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+
+
+def _env_qdrant_api_key() -> Optional[str]:
+    return (
+        os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+
+
 def load_config(
     config_path: Optional[str] = None,
     force_reload: bool = False,
@@ -171,8 +184,8 @@ def get_section(section: str) -> Dict[str, Any]:
 def get_qdrant_config() -> Dict[str, Any]:
     """Get Qdrant configuration with defaults."""
     return {
-        "url": get("qdrant.url", os.getenv("QDRANT_URL")),
-        "api_key": get("qdrant.api_key", os.getenv("QDRANT_API_KEY")),
+        "url": get("qdrant.url", _env_qdrant_url()),
+        "api_key": get("qdrant.api_key", _env_qdrant_api_key()),
         "collection": get("qdrant.collection", "visual_documents"),
     }
 

visual_rag/embedding/pooling.py

@@ -15,10 +15,27 @@ import logging
 logger = logging.getLogger(__name__)
 
 
+def _infer_output_dtype(
+    embedding: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
+) -> np.dtype:
+    """Infer output dtype: use provided, else match input (fp16→fp16, bf16→fp32, fp32→fp32)."""
+    if output_dtype is not None:
+        return output_dtype
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.float16:
+            return np.float16
+        return np.float32
+    if isinstance(embedding, np.ndarray) and embedding.dtype == np.float16:
+        return np.float16
+    return np.float32
+
+
 def tile_level_mean_pooling(
     embedding: Union[torch.Tensor, np.ndarray],
     num_tiles: int,
     patches_per_tile: int = 64,
+    output_dtype: Optional[np.dtype] = None,
 ) -> np.ndarray:
     """
     Compute tile-level mean pooling for multi-vector embeddings.
@@ -35,6 +52,7 @@ def tile_level_mean_pooling(
         embedding: Visual token embeddings [num_visual_tokens, dim]
         num_tiles: Number of tiles (including global tile)
         patches_per_tile: Patches per tile (64 for ColSmol)
+        output_dtype: Output dtype (default: infer from input, fp16→fp16, bf16→fp32)
     
     Returns:
         Tile-level pooled embeddings [num_tiles, dim]
@@ -45,7 +63,7 @@ def tile_level_mean_pooling(
         >>> pooled = tile_level_mean_pooling(embedding, num_tiles=13)
         >>> print(pooled.shape)  # (13, 128)
     """
-    # Convert to numpy
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
     if isinstance(embedding, torch.Tensor):
         if embedding.dtype == torch.bfloat16:
             emb_np = embedding.cpu().float().numpy()
@@ -57,17 +75,15 @@ def tile_level_mean_pooling(
     num_visual_tokens = emb_np.shape[0]
     expected_tokens = num_tiles * patches_per_tile
     
-    # Handle mismatch (can happen with global tile variations)
     if num_visual_tokens != expected_tokens:
         logger.debug(
             f"Token count mismatch: {num_visual_tokens} vs expected {expected_tokens}"
         )
         actual_tiles = num_visual_tokens // patches_per_tile
         if actual_tiles * patches_per_tile != num_visual_tokens:
-            actual_tiles += 1  # Include partial tile
+            actual_tiles += 1
         num_tiles = actual_tiles
     
-    # Compute mean per tile
     tile_embeddings = []
     for tile_idx in range(num_tiles):
         start_idx = tile_idx * patches_per_tile
@@ -80,11 +96,134 @@ def tile_level_mean_pooling(
         tile_mean = tile_patches.mean(axis=0)
         tile_embeddings.append(tile_mean)
     
-    return np.array(tile_embeddings, dtype=np.float32)
+    return np.array(tile_embeddings, dtype=out_dtype)
+
+
+def colpali_row_mean_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    grid_size: int = 32,
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy().astype(np.float32)
+    else:
+        emb_np = np.array(embedding, dtype=np.float32)
+
+    num_tokens, dim = emb_np.shape
+    expected = int(grid_size) * int(grid_size)
+    if num_tokens != expected:
+        raise ValueError(f"Expected {expected} visual tokens for grid_size={grid_size}, got {num_tokens}")
+
+    grid = emb_np.reshape(int(grid_size), int(grid_size), int(dim))
+    pooled = grid.mean(axis=1)
+    return pooled.astype(out_dtype)
+
+
+def colsmol_experimental_pooling(
+    embedding: Union[torch.Tensor, np.ndarray],
+    num_tiles: int,
+    patches_per_tile: int = 64,
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
+    if isinstance(embedding, torch.Tensor):
+        if embedding.dtype == torch.bfloat16:
+            emb_np = embedding.cpu().float().numpy()
+        else:
+            emb_np = embedding.cpu().numpy().astype(np.float32)
+    else:
+        emb_np = np.array(embedding, dtype=np.float32)
+
+    num_visual_tokens, dim = emb_np.shape
+    if num_tiles <= 0:
+        raise ValueError("num_tiles must be > 0")
+    if patches_per_tile <= 0:
+        raise ValueError("patches_per_tile must be > 0")
+
+    last_tile_start = (int(num_tiles) - 1) * int(patches_per_tile)
+    if last_tile_start >= num_visual_tokens:
+        actual_tiles = int(num_visual_tokens) // int(patches_per_tile)
+        if actual_tiles * int(patches_per_tile) != int(num_visual_tokens):
+            actual_tiles += 1
+        if actual_tiles <= 0:
+            raise ValueError(
+                f"Not enough tokens for num_tiles={num_tiles}, patches_per_tile={patches_per_tile}: got {num_visual_tokens}"
+            )
+        num_tiles = actual_tiles
+        last_tile_start = (int(num_tiles) - 1) * int(patches_per_tile)
+
+    prefix = emb_np[:last_tile_start]
+    last_tile = emb_np[last_tile_start : min(last_tile_start + int(patches_per_tile), num_visual_tokens)]
+
+    if prefix.size:
+        prefix_tiles = prefix.reshape(-1, int(patches_per_tile), int(dim))
+        prefix_means = prefix_tiles.mean(axis=1)
+    else:
+        prefix_means = np.zeros((0, int(dim)), dtype=out_dtype)
+
+    return np.concatenate([prefix_means.astype(out_dtype), last_tile.astype(out_dtype)], axis=0)
+
+
+def colpali_experimental_pooling_from_rows(
+    row_vectors: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
+) -> np.ndarray:
+    """
+    Experimental "convolution-style" pooling with window size 3.
+    
+    For N input rows, produces N + 2 output vectors:
+    - Position 0: row[0] alone (1 row)
+    - Position 1: mean(rows[0:2]) (2 rows)
+    - Position 2: mean(rows[0:3]) (3 rows)
+    - Positions 3 to N-1: sliding window of 3 (rows[i-2:i+1])
+    - Position N: mean(rows[N-2:N]) (last 2 rows)
+    - Position N+1: row[N-1] alone (last row)
+    
+    For N=32 rows: produces 34 vectors.
+    """
+    out_dtype = _infer_output_dtype(row_vectors, output_dtype)
+    if isinstance(row_vectors, torch.Tensor):
+        if row_vectors.dtype == torch.bfloat16:
+            rows = row_vectors.cpu().float().numpy()
+        else:
+            rows = row_vectors.cpu().numpy().astype(np.float32)
+    else:
+        rows = np.array(row_vectors, dtype=np.float32)
+
+    n, dim = rows.shape
+    if n < 1:
+        raise ValueError("row_vectors must be non-empty")
+    if n == 1:
+        return rows.astype(out_dtype)
+    if n == 2:
+        return np.stack([rows[0], rows[:2].mean(axis=0), rows[1]], axis=0).astype(out_dtype)
+    if n == 3:
+        return np.stack([
+            rows[0],
+            rows[:2].mean(axis=0),
+            rows[:3].mean(axis=0),
+            rows[1:3].mean(axis=0),
+            rows[2],
+        ], axis=0).astype(out_dtype)
+
+    out = np.zeros((n + 2, dim), dtype=np.float32)
+    out[0] = rows[0]
+    out[1] = rows[:2].mean(axis=0)
+    out[2] = rows[:3].mean(axis=0)
+    for i in range(3, n):
+        out[i] = rows[i - 2 : i + 1].mean(axis=0)
+    out[n] = rows[n - 2 : n].mean(axis=0)
+    out[n + 1] = rows[n - 1]
+    return out.astype(out_dtype)
 
 
 def global_mean_pooling(
     embedding: Union[torch.Tensor, np.ndarray],
+    output_dtype: Optional[np.dtype] = None,
 ) -> np.ndarray:
     """
     Compute global mean pooling → single vector.
@@ -94,10 +233,12 @@ def global_mean_pooling(
     
     Args:
         embedding: Multi-vector embeddings [num_tokens, dim]
+        output_dtype: Output dtype (default: infer from input, fp16→fp16, bf16→fp32)
     
     Returns:
         Pooled vector [dim]
     """
+    out_dtype = _infer_output_dtype(embedding, output_dtype)
     if isinstance(embedding, torch.Tensor):
         if embedding.dtype == torch.bfloat16:
             emb_np = embedding.cpu().float().numpy()
@@ -106,7 +247,7 @@ def global_mean_pooling(
     else:
         emb_np = np.array(embedding)
     
-    return emb_np.mean(axis=0).astype(np.float32)
+    return emb_np.mean(axis=0).astype(out_dtype)
 
 
 def compute_maxsim_score(
@@ -198,5 +339,3 @@ def compute_maxsim_batch(
         scores.append(float(max_sims.sum()))
     
     return scores
-
-

visual_rag/embedding/visual_embedder.py

@@ -70,24 +70,26 @@ class VisualEmbedder:
         backend: str = "auto",
         device: Optional[str] = None,
         torch_dtype: Optional[torch.dtype] = None,
+        output_dtype: Optional[np.dtype] = None,
         batch_size: int = 4,
         filter_special_tokens: bool = True,
+        processor_speed: str = "fast",
     ):
         self.model_name = model_name
         self.batch_size = batch_size
         self.filter_special_tokens = filter_special_tokens
+        if processor_speed not in ("fast", "slow", "auto"):
+            raise ValueError("processor_speed must be one of: fast, slow, auto")
+        self.processor_speed = processor_speed
         
-        # Allow override via environment variable
         if os.getenv("VISUALRAG_INCLUDE_SPECIAL_TOKENS"):
             self.filter_special_tokens = False
             logger.info("Special token filtering disabled via VISUALRAG_INCLUDE_SPECIAL_TOKENS")
         
-        # Auto-detect backend from model name
         if backend == "auto":
             backend = self._detect_backend(model_name)
         self.backend = backend
         
-        # Auto-detect device
         if device is None:
             if torch.cuda.is_available():
                 device = "cuda"
@@ -97,7 +99,6 @@ class VisualEmbedder:
                 device = "cpu"
         self.device = device
         
-        # Auto-detect dtype
         if torch_dtype is None:
             if device == "cuda":
                 torch_dtype = torch.bfloat16
@@ -105,7 +106,13 @@ class VisualEmbedder:
                 torch_dtype = torch.float32
         self.torch_dtype = torch_dtype
         
-        # Lazy model loading
+        if output_dtype is None:
+            if torch_dtype == torch.float16:
+                output_dtype = np.float16
+            else:
+                output_dtype = np.float32
+        self.output_dtype = output_dtype
+        
         self._model = None
         self._processor = None
         self._image_token_id = None
@@ -113,7 +120,7 @@ class VisualEmbedder:
         logger.info(f"🤖 VisualEmbedder initialized")
         logger.info(f"   Model: {model_name}")
         logger.info(f"   Backend: {backend}")
-        logger.info(f"   Device: {device}, dtype: {torch_dtype}")
+        logger.info(f"   Device: {device}, torch_dtype: {torch_dtype}, output_dtype: {output_dtype}")
     
     def _detect_backend(self, model_name: str) -> str:
         """Auto-detect backend from model name."""
@@ -141,7 +148,14 @@ class VisualEmbedder:
     def _load_colpali_model(self):
         """Load ColPali-family model."""
         try:
-            from colpali_engine.models import ColIdefics3, ColIdefics3Processor
+            from colpali_engine.models import (
+                ColIdefics3,
+                ColIdefics3Processor,
+                ColPali,
+                ColPaliProcessor,
+                ColQwen2,
+                ColQwen2Processor,
+            )
         except ImportError:
             raise ImportError(
                 "colpali_engine not installed. Install with: "
@@ -151,25 +165,80 @@ class VisualEmbedder:
         
         logger.info(f"🤖 Loading ColPali model: {self.model_name}")
         logger.info(f"   Device: {self.device}, dtype: {self.torch_dtype}")
+
+        def _processor_kwargs():
+            if self.processor_speed == "auto":
+                return {}
+            return {"use_fast": self.processor_speed == "fast"}
         
-        # Determine attention implementation
+        from transformers import AutoConfig
+
+        cfg = AutoConfig.from_pretrained(self.model_name)
+        model_type = str(getattr(cfg, "model_type", "") or "").lower()
+
+        if model_type == "colpali" or "colpali" in (self.model_name or "").lower():
+            self._model = ColPali.from_pretrained(
+                self.model_name,
+                torch_dtype=self.torch_dtype,
+                device_map=self.device,
+            ).eval()
+            try:
+                self._processor = ColPaliProcessor.from_pretrained(self.model_name, **_processor_kwargs())
+            except TypeError:
+                self._processor = ColPaliProcessor.from_pretrained(self.model_name)
+            except Exception:
+                if self.processor_speed == "fast":
+                    self._processor = ColPaliProcessor.from_pretrained(self.model_name, use_fast=False)
+                else:
+                    raise
+            self._image_token_id = self._processor.image_token_id
+            logger.info("✅ Loaded ColPali backend")
+            return
+
+        if model_type.startswith("qwen2") or "colqwen" in (self.model_name or "").lower():
+            self._model = ColQwen2.from_pretrained(
+                self.model_name,
+                dtype=self.torch_dtype,
+                device_map=self.device,
+            ).eval()
+            try:
+                self._processor = ColQwen2Processor.from_pretrained(self.model_name, device_map=self.device, **_processor_kwargs())
+            except TypeError:
+                self._processor = ColQwen2Processor.from_pretrained(self.model_name, device_map=self.device)
+            except Exception:
+                if self.processor_speed == "fast":
+                    self._processor = ColQwen2Processor.from_pretrained(self.model_name, device_map=self.device, use_fast=False)
+                else:
+                    raise
+            self._image_token_id = self._processor.image_token_id
+            logger.info("✅ Loaded ColQwen2 backend")
+            return
+
         attn_implementation = "eager"
         if self.device != "cpu":
             try:
                 import flash_attn  # noqa
+
                 attn_implementation = "flash_attention_2"
                 logger.info("   Using FlashAttention2")
             except ImportError:
                 pass
-        
+
         self._model = ColIdefics3.from_pretrained(
             self.model_name,
             dtype=self.torch_dtype,
             device_map=self.device,
             attn_implementation=attn_implementation,
         ).eval()
-        
-        self._processor = ColIdefics3Processor.from_pretrained(self.model_name)
+        try:
+            self._processor = ColIdefics3Processor.from_pretrained(self.model_name, **_processor_kwargs())
+        except TypeError:
+            self._processor = ColIdefics3Processor.from_pretrained(self.model_name)
+        except Exception:
+            if self.processor_speed == "fast":
+                self._processor = ColIdefics3Processor.from_pretrained(self.model_name, use_fast=False)
+            else:
+                raise
         self._image_token_id = self._processor.image_token_id
         
         logger.info("✅ Model loaded successfully")
@@ -244,6 +313,61 @@ class VisualEmbedder:
             logger.debug(f"Query embedding: {embedding.shape[0]} tokens (unfiltered)")
         
         return embedding
+
+    def embed_queries(
+        self,
+        query_texts: List[str],
+        batch_size: Optional[int] = None,
+        filter_special_tokens: Optional[bool] = None,
+        show_progress: bool = True,
+    ) -> List[torch.Tensor]:
+        """
+        Generate embeddings for a list of text queries.
+
+        Returns a list of tensors, each of shape [num_tokens, embedding_dim].
+        """
+        should_filter = (
+            filter_special_tokens if filter_special_tokens is not None else self.filter_special_tokens
+        )
+        batch_size = batch_size or self.batch_size
+
+        outputs: List[torch.Tensor] = []
+        iterator = range(0, len(query_texts), batch_size)
+        if show_progress:
+            iterator = tqdm(iterator, desc="📝 Embedding queries", unit="batch")
+
+        for i in iterator:
+            batch = query_texts[i : i + batch_size]
+            with torch.no_grad():
+                processed = self.processor.process_queries(batch).to(self.model.device)
+                batch_embeddings = self.model(**processed)
+
+            if isinstance(batch_embeddings, torch.Tensor) and batch_embeddings.dim() == 3:
+                attn = processed.get("attention_mask") if should_filter else None
+                input_ids = processed.get("input_ids") if should_filter else None
+
+                for j in range(batch_embeddings.shape[0]):
+                    emb = batch_embeddings[j]
+                    if should_filter and attn is not None:
+                        valid_mask = attn[j].bool()
+                        emb = emb[valid_mask]
+                        if input_ids is not None:
+                            ids = input_ids[j][valid_mask]
+                            non_special_mask = ids >= 4
+                            if non_special_mask.any():
+                                emb = emb[non_special_mask]
+                    outputs.append(emb)
+            else:
+                outputs.extend(batch_embeddings)
+
+            del processed, batch_embeddings
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            elif torch.backends.mps.is_available():
+                torch.mps.empty_cache()
+
+        return outputs
     
     def embed_images(
         self,
@@ -274,6 +398,8 @@ class VisualEmbedder:
             - num_tiles: Total tiles (n_rows × n_cols + 1 global)
         """
         batch_size = batch_size or self.batch_size
+        if self.device == "mps" and "colpali" in (self.model_name or "").lower() and int(batch_size) > 1:
+            batch_size = 1
         
         embeddings = []
         token_infos = [] if return_token_info else None
@@ -325,6 +451,8 @@ class VisualEmbedder:
             gc.collect()
             if torch.cuda.is_available():
                 torch.cuda.empty_cache()
+            elif torch.backends.mps.is_available():
+                torch.mps.empty_cache()
         
         if return_token_info:
             return embeddings, token_infos
@@ -357,10 +485,92 @@ class VisualEmbedder:
         else:
             visual_emb = np.array(full_embedding)[visual_indices]
         
-        return visual_emb.astype(np.float32)
+        return visual_emb.astype(self.output_dtype)
 
+    def mean_pool_visual_embedding(
+        self,
+        visual_embedding: Union[torch.Tensor, np.ndarray],
+        token_info: Optional[Dict[str, Any]] = None,
+        *,
+        target_vectors: int = 32,
+    ) -> np.ndarray:
+        from visual_rag.embedding.pooling import colpali_row_mean_pooling, tile_level_mean_pooling
 
-# Backward compatibility alias
-ColPaliEmbedder = VisualEmbedder
+        model_lower = (self.model_name or "").lower()
+        is_colsmol = "colsmol" in model_lower
 
+        if isinstance(visual_embedding, torch.Tensor):
+            if visual_embedding.dtype == torch.bfloat16:
+                visual_np = visual_embedding.cpu().float().numpy()
+            else:
+                visual_np = visual_embedding.cpu().numpy().astype(np.float32)
+        else:
+            visual_np = np.array(visual_embedding, dtype=np.float32)
 
+        if is_colsmol:
+            n_rows = (token_info or {}).get("n_rows")
+            n_cols = (token_info or {}).get("n_cols")
+            num_tiles = int(n_rows) * int(n_cols) + 1 if n_rows and n_cols else 13
+            return tile_level_mean_pooling(visual_np, num_tiles=num_tiles, patches_per_tile=64, output_dtype=self.output_dtype)
+
+        num_tokens = int(visual_np.shape[0])
+        grid = int(round(float(num_tokens) ** 0.5))
+        if grid * grid != num_tokens:
+            raise ValueError(f"Cannot infer square grid from num_visual_tokens={num_tokens} for model={self.model_name}")
+        if int(target_vectors) != int(grid):
+            raise ValueError(
+                f"target_vectors={target_vectors} does not match inferred grid_size={grid} for model={self.model_name}"
+            )
+        return colpali_row_mean_pooling(visual_np, grid_size=int(target_vectors), output_dtype=self.output_dtype)
+
+    def global_pool_from_mean_pool(self, mean_pool: np.ndarray) -> np.ndarray:
+        if mean_pool.size == 0:
+            return np.zeros((128,), dtype=self.output_dtype)
+        return mean_pool.mean(axis=0).astype(self.output_dtype)
+
+    def experimental_pool_visual_embedding(
+        self,
+        visual_embedding: Union[torch.Tensor, np.ndarray],
+        token_info: Optional[Dict[str, Any]] = None,
+        *,
+        target_vectors: int = 32,
+        mean_pool: Optional[np.ndarray] = None,
+    ) -> np.ndarray:
+        from visual_rag.embedding.pooling import colpali_experimental_pooling_from_rows, colsmol_experimental_pooling
+
+        model_lower = (self.model_name or "").lower()
+        is_colsmol = "colsmol" in model_lower
+
+        if isinstance(visual_embedding, torch.Tensor):
+            if visual_embedding.dtype == torch.bfloat16:
+                visual_np = visual_embedding.cpu().float().numpy()
+            else:
+                visual_np = visual_embedding.cpu().numpy().astype(np.float32)
+        else:
+            visual_np = np.array(visual_embedding, dtype=np.float32)
+
+        if is_colsmol:
+            if mean_pool is not None and getattr(mean_pool, "shape", None) is not None and int(mean_pool.shape[0]) > 0:
+                num_tiles = int(mean_pool.shape[0])
+            else:
+                num_tiles = (token_info or {}).get("num_tiles")
+                if num_tiles is None:
+                    num_visual_tokens = (token_info or {}).get("num_visual_tokens")
+                    if num_visual_tokens is None:
+                        num_visual_tokens = int(visual_np.shape[0])
+                    patches_per_tile = 64
+                    num_tiles = int(num_visual_tokens) // patches_per_tile
+                    if int(num_tiles) * patches_per_tile != int(num_visual_tokens):
+                        num_tiles = int(num_tiles) + 1
+                num_tiles = int(num_tiles)
+            return colsmol_experimental_pooling(visual_np, num_tiles=num_tiles, patches_per_tile=64, output_dtype=self.output_dtype)
+
+        rows = mean_pool if mean_pool is not None else self.mean_pool_visual_embedding(visual_np, token_info, target_vectors=target_vectors)
+        if int(rows.shape[0]) != int(target_vectors):
+            raise ValueError(
+                f"experimental pooling expects mean_pool to have {target_vectors} rows, got {rows.shape[0]} for model={self.model_name}"
+            )
+        return colpali_experimental_pooling_from_rows(rows, output_dtype=self.output_dtype)
+
+# Backward compatibility alias
+ColPaliEmbedder = VisualEmbedder

visual_rag/indexing/cloudinary_uploader.py

@@ -8,6 +8,10 @@ Features:
 - Retry logic with timeouts
 - Batch uploading
 - Automatic JPEG optimization
+
+Environment Variables:
+- VISUAL_RAG_THREAD_SAFE: Set to "1" to use thread-safe timeouts
+  (required for Streamlit, Flask, or other threaded contexts)
 """
 
 import io
@@ -16,12 +20,16 @@ import time
 import signal
 import logging
 import platform
+import threading
 from typing import Optional
+from concurrent.futures import ThreadPoolExecutor, TimeoutError as FuturesTimeoutError
 
 from PIL import Image
 
 logger = logging.getLogger(__name__)
 
+THREAD_SAFE_MODE = os.getenv("VISUAL_RAG_THREAD_SAFE", "").lower() in ("1", "true", "yes")
+
 
 class CloudinaryUploader:
     """
@@ -124,40 +132,74 @@ class CloudinaryUploader:
         if subfolder:
             folder_path = f"{self.folder}/{subfolder}"
         
-        # Timeout handling (Unix/macOS only)
+        def do_upload():
+            buffer.seek(0)
+            result = cloudinary.uploader.upload(
+                buffer,
+                folder=folder_path,
+                overwrite=True,
+                public_id=public_id,
+                resource_type="image",
+                timeout=self.timeout_seconds,
+            )
+            return result["secure_url"]
+        
+        # Use thread-safe mode for Streamlit/Flask/threaded contexts
+        # Set VISUAL_RAG_THREAD_SAFE=1 to enable
+        if THREAD_SAFE_MODE or threading.current_thread() is not threading.main_thread():
+            return self._upload_with_thread_timeout(do_upload, public_id)
+        else:
+            return self._upload_with_signal_timeout(do_upload, public_id)
+    
+    def _upload_with_thread_timeout(self, do_upload, public_id: str) -> Optional[str]:
+        """Thread-safe upload with ThreadPoolExecutor timeout."""
+        for attempt in range(self.max_retries):
+            try:
+                with ThreadPoolExecutor(max_workers=1) as executor:
+                    future = executor.submit(do_upload)
+                    return future.result(timeout=self.timeout_seconds)
+                        
+            except FuturesTimeoutError:
+                logger.warning(
+                    f"Upload timeout (attempt {attempt + 1}/{self.max_retries}): {public_id}"
+                )
+                if attempt < self.max_retries - 1:
+                    time.sleep(2 ** attempt)
+                    
+            except Exception as e:
+                logger.warning(
+                    f"Upload failed (attempt {attempt + 1}/{self.max_retries}): {e}"
+                )
+                if attempt < self.max_retries - 1:
+                    time.sleep(2 ** attempt)
+        
+        logger.error(f"❌ Upload failed after {self.max_retries} attempts: {public_id}")
+        return None
+    
+    def _upload_with_signal_timeout(self, do_upload, public_id: str) -> Optional[str]:
+        """Signal-based upload timeout (main thread only, Unix/macOS)."""
         use_timeout = platform.system() != "Windows"
         
-        class TimeoutError(Exception):
+        class SignalTimeoutError(Exception):
             pass
         
         def timeout_handler(signum, frame):
-            raise TimeoutError(f"Upload timed out after {self.timeout_seconds}s")
+            raise SignalTimeoutError(f"Upload timed out after {self.timeout_seconds}s")
         
         for attempt in range(self.max_retries):
             try:
-                buffer.seek(0)
-                
-                # Set timeout alarm
                 if use_timeout:
                     old_handler = signal.signal(signal.SIGALRM, timeout_handler)
                     signal.alarm(self.timeout_seconds)
                 
                 try:
-                    result = cloudinary.uploader.upload(
-                        buffer,
-                        folder=folder_path,
-                        overwrite=True,
-                        public_id=public_id,
-                        resource_type="image",
-                        timeout=self.timeout_seconds,
-                    )
-                    return result["secure_url"]
+                    return do_upload()
                 finally:
                     if use_timeout:
                         signal.alarm(0)
                         signal.signal(signal.SIGALRM, old_handler)
                         
-            except TimeoutError:
+            except SignalTimeoutError:
                 logger.warning(
                     f"Upload timeout (attempt {attempt + 1}/{self.max_retries}): {public_id}"
                 )
@@ -205,4 +247,33 @@ class CloudinaryUploader:
         
         return original_url, resized_url
 
+    def upload_original_cropped_and_resized(
+        self,
+        original_image: Image.Image,
+        cropped_image: Optional[Image.Image],
+        resized_image: Image.Image,
+        base_public_id: str,
+    ) -> tuple:
+        original_url = self.upload(
+            original_image,
+            base_public_id,
+            subfolder="original",
+        )
+
+        cropped_url = None
+        if cropped_image is not None:
+            cropped_url = self.upload(
+                cropped_image,
+                base_public_id,
+                subfolder="cropped",
+            )
+
+        resized_url = self.upload(
+            resized_image,
+            base_public_id,
+            subfolder="resized",
+        )
+
+        return original_url, cropped_url, resized_url
+
 

visual_rag/indexing/pipeline.py

@@ -89,6 +89,11 @@ class ProcessingPipeline:
         metadata_mapping: Optional[Dict[str, Dict[str, Any]]] = None,
         config: Optional[Dict[str, Any]] = None,
         embedding_strategy: str = "pooling",
+        crop_empty: bool = False,
+        crop_empty_percentage_to_remove: float = 0.9,
+        crop_empty_remove_page_number: bool = False,
+        crop_empty_preserve_border_px: int = 1,
+        crop_empty_uniform_rowcol_std_threshold: float = 0.0,
     ):
         self.embedder = embedder
         self.indexer = indexer
@@ -103,6 +108,12 @@ class ProcessingPipeline:
                 f"Must be one of: {self.STRATEGIES}"
             )
         self.embedding_strategy = embedding_strategy
+
+        self.crop_empty = bool(crop_empty)
+        self.crop_empty_percentage_to_remove = float(crop_empty_percentage_to_remove)
+        self.crop_empty_remove_page_number = bool(crop_empty_remove_page_number)
+        self.crop_empty_preserve_border_px = int(crop_empty_preserve_border_px)
+        self.crop_empty_uniform_rowcol_std_threshold = float(crop_empty_uniform_rowcol_std_threshold)
         
         logger.info(f"📊 Embedding strategy: {embedding_strategy}")
         if embedding_strategy == "pooling":
@@ -128,6 +139,8 @@ class ProcessingPipeline:
         skip_existing: bool = True,
         upload_to_cloudinary: bool = True,
         upload_to_qdrant: bool = True,
+        original_filename: Optional[str] = None,
+        progress_callback: Optional[callable] = None,
     ) -> Dict[str, Any]:
         """
         Process a single PDF end-to-end.
@@ -137,6 +150,8 @@ class ProcessingPipeline:
             skip_existing: Skip pages that already exist in Qdrant
             upload_to_cloudinary: Upload images to Cloudinary
             upload_to_qdrant: Upload embeddings to Qdrant
+            original_filename: Original filename (use this instead of pdf_path.name for temp files)
+            progress_callback: Optional callback(stage, current, total, message) for progress updates
         
         Returns:
             Dict with processing results:
@@ -150,23 +165,26 @@ class ProcessingPipeline:
             }
         """
         pdf_path = Path(pdf_path)
-        logger.info(f"📚 Processing PDF: {pdf_path.name}")
+        filename = original_filename or pdf_path.name
+        logger.info(f"📚 Processing PDF: {filename}")
         
         # Check existing pages
         existing_ids: Set[str] = set()
         if skip_existing and self.indexer:
-            existing_ids = self.indexer.get_existing_ids(pdf_path.name)
+            existing_ids = self.indexer.get_existing_ids(filename)
             if existing_ids:
                 logger.info(f"   Found {len(existing_ids)} existing pages")
         
-        # Convert PDF to images
         logger.info(f"🖼️ Converting PDF to images...")
+        if progress_callback:
+            progress_callback("convert", 0, 0, "Converting PDF to images...")
         images, texts = self.pdf_processor.process_pdf(pdf_path)
         total_pages = len(images)
         logger.info(f"   ✅ Converted {total_pages} pages")
+        if progress_callback:
+            progress_callback("convert", total_pages, total_pages, f"Converted {total_pages} pages")
         
-        # Get extra metadata
-        extra_metadata = self._get_extra_metadata(pdf_path.name)
+        extra_metadata = self._get_extra_metadata(filename)
         if extra_metadata:
             logger.info(f"   📋 Found extra metadata: {list(extra_metadata.keys())}")
         
@@ -183,12 +201,13 @@ class ProcessingPipeline:
             batch_texts = texts[batch_start:batch_end]
             
             logger.info(f"📦 Processing pages {batch_start + 1}-{batch_end}/{total_pages}")
+            if progress_callback:
+                progress_callback("embed", batch_start, total_pages, f"Embedding pages {batch_start + 1}-{batch_end}")
             
-            # Filter pages that need processing
             pages_to_process = []
             for i, (img, text) in enumerate(zip(batch_images, batch_texts)):
                 page_num = batch_start + i + 1
-                chunk_id = self.generate_chunk_id(pdf_path.name, page_num)
+                chunk_id = self.generate_chunk_id(filename, page_num)
                 
                 if skip_existing and chunk_id in existing_ids:
                     skipped += 1
@@ -198,7 +217,7 @@ class ProcessingPipeline:
                     "index": i,
                     "page_num": page_num,
                     "chunk_id": chunk_id,
-                    "image": img,
+                    "raw_image": img,
                     "text": text,
                 })
             
@@ -208,36 +227,64 @@ class ProcessingPipeline:
             
             # Generate embeddings with token info
             logger.info(f"🤖 Generating embeddings for {len(pages_to_process)} pages...")
-            images_to_embed = [p["image"] for p in pages_to_process]
+            from visual_rag.preprocessing.crop_empty import CropEmptyConfig, crop_empty
+
+            images_to_embed = []
+            for p in pages_to_process:
+                raw_img = p["raw_image"]
+                if self.crop_empty:
+                    cropped_img, crop_meta = crop_empty(
+                        raw_img,
+                        config=CropEmptyConfig(
+                            percentage_to_remove=float(self.crop_empty_percentage_to_remove),
+                            remove_page_number=bool(self.crop_empty_remove_page_number),
+                            preserve_border_px=int(self.crop_empty_preserve_border_px),
+                            uniform_rowcol_std_threshold=float(self.crop_empty_uniform_rowcol_std_threshold),
+                        ),
+                    )
+                    p["embed_image"] = cropped_img
+                    p["crop_meta"] = crop_meta
+                    images_to_embed.append(cropped_img)
+                else:
+                    p["embed_image"] = raw_img
+                    p["crop_meta"] = None
+                    images_to_embed.append(raw_img)
             
             embeddings, token_infos = self.embedder.embed_images(
                 images_to_embed,
                 batch_size=self.embedding_batch_size,
                 return_token_info=True,
-                show_progress=True,
+                show_progress=False,
             )
             
-            # Process each page
             for idx, page_info in enumerate(pages_to_process):
-                orig_img = page_info["image"]
+                raw_img = page_info["raw_image"]
+                embed_img = page_info["embed_image"]
+                crop_meta = page_info["crop_meta"]
                 page_num = page_info["page_num"]
                 chunk_id = page_info["chunk_id"]
                 text = page_info["text"]
                 embedding = embeddings[idx]
                 token_info = token_infos[idx]
                 
+                if progress_callback:
+                    progress_callback("process", page_num, total_pages, f"Processing page {page_num}/{total_pages}")
+                
                 try:
                     page_data = self._process_single_page(
-                        pdf_path=pdf_path,
+                        filename=filename,
+                        pdf_stem=pdf_path.stem,
                         page_num=page_num,
                         chunk_id=chunk_id,
                         total_pages=total_pages,
-                        orig_img=orig_img,
+                        raw_img=raw_img,
+                        embed_img=embed_img,
                         text=text,
                         embedding=embedding,
                         token_info=token_info,
                         extra_metadata=extra_metadata,
                         upload_to_cloudinary=upload_to_cloudinary,
+                        crop_meta=crop_meta,
                     )
                     
                     all_pages.append(page_data)
@@ -265,10 +312,10 @@ class ProcessingPipeline:
             count = self._upload_batch(upload_queue)
             uploaded += count
         
-        logger.info(f"✅ Completed {pdf_path.name}: {uploaded} uploaded, {skipped} skipped, {failed} failed")
+        logger.info(f"✅ Completed {filename}: {uploaded} uploaded, {skipped} skipped, {failed} failed")
         
         return {
-            "filename": pdf_path.name,
+            "filename": filename,
             "total_pages": total_pages,
             "uploaded": uploaded,
             "skipped": skipped,
@@ -278,22 +325,25 @@ class ProcessingPipeline:
     
     def _process_single_page(
         self,
-        pdf_path: Path,
+        filename: str,
+        pdf_stem: str,
         page_num: int,
         chunk_id: str,
         total_pages: int,
-        orig_img,
+        raw_img,
+        embed_img,
         text: str,
         embedding: torch.Tensor,
         token_info: Dict[str, Any],
         extra_metadata: Dict[str, Any],
         upload_to_cloudinary: bool = True,
+        crop_meta: Optional[Dict[str, Any]] = None,
     ) -> Dict[str, Any]:
         """Process a single page with full metadata for saliency."""
-        from visual_rag.embedding.pooling import tile_level_mean_pooling, global_mean_pooling
+        from visual_rag.embedding.pooling import global_mean_pooling
         
         # Resize image for ColPali
-        resized_img, tile_rows, tile_cols = self.pdf_processor.resize_for_colpali(orig_img)
+        resized_img, tile_rows, tile_cols = self.pdf_processor.resize_for_colpali(embed_img)
         
         # Use processor's tile info if available (more accurate)
         proc_n_rows = token_info.get("n_rows")
@@ -316,10 +366,6 @@ class ProcessingPipeline:
         visual_indices = token_info["visual_token_indices"]
         num_visual_tokens = token_info["num_visual_tokens"]
         
-        # Compute tile structure
-        num_tiles = tile_rows * tile_cols + 1  # +1 for global tile
-        patches_per_tile = 64
-        
         # =========================================================================
         # STRATEGY: "pooling" (NOVEL) vs "standard" (BASELINE) vs "all" (BOTH)
         # =========================================================================
@@ -327,27 +373,39 @@ class ProcessingPipeline:
         # Always compute visual-only embedding (needed for pooling and saliency)
         visual_embedding = full_embedding[visual_indices]
         
-        # Compute tile-level mean pooling: [num_visual_tokens, 128] → [num_tiles, 128]
-        tile_pooled = tile_level_mean_pooling(visual_embedding, num_tiles, patches_per_tile)
-        
-        # Compute global mean (baseline pooling)
+        tile_pooled = self.embedder.mean_pool_visual_embedding(visual_embedding, token_info, target_vectors=32)
+        experimental_pooled = self.embedder.experimental_pool_visual_embedding(
+            visual_embedding, token_info, target_vectors=32, mean_pool=tile_pooled
+        )
         global_pooled = global_mean_pooling(full_embedding)
+        global_pooling = self.embedder.global_pool_from_mean_pool(tile_pooled) if tile_pooled.size else global_pooled
+
+        num_tiles = int(tile_pooled.shape[0])
+        patches_per_tile = int(visual_embedding.shape[0] // max(num_tiles, 1)) if num_tiles else 0
+        if tile_rows and tile_cols and int(tile_rows) * int(tile_cols) + 1 == num_tiles:
+            pass
+        else:
+            tile_rows = token_info.get("n_rows") or None
+            tile_cols = token_info.get("n_cols") or None
         
         if self.embedding_strategy == "pooling":
             # NOVEL APPROACH: Visual tokens only + tile-level pooling
             embedding_for_initial = visual_embedding
             embedding_for_pooling = tile_pooled
+            global_pooling = self.embedder.global_pool_from_mean_pool(tile_pooled) if tile_pooled.size else global_pooled
             
         elif self.embedding_strategy == "standard":
             # BASELINE: All tokens + global mean
             embedding_for_initial = full_embedding
             embedding_for_pooling = global_pooled.reshape(1, -1)
+            global_pooling = global_pooled
             
         else:  # "all" - Push BOTH representations (efficient for comparison)
             # Embed once, store multiple vector representations
             # This allows comparing both strategies without re-embedding
             embedding_for_initial = visual_embedding  # Use visual for search
             embedding_for_pooling = tile_pooled       # Use tile-level for fast prefetch
+            global_pooling = self.embedder.global_pool_from_mean_pool(tile_pooled) if tile_pooled.size else global_pooled
             
             # ALSO store standard representations as additional vectors
             # These will be added to metadata for optional use
@@ -355,21 +413,25 @@ class ProcessingPipeline:
         
         # Upload to Cloudinary
         original_url = None
+        cropped_url = None
         resized_url = None
         
         if upload_to_cloudinary and self.cloudinary_uploader:
-            base_filename = f"{pdf_path.stem}_page_{page_num}"
-            original_url, resized_url = self.cloudinary_uploader.upload_original_and_resized(
-                orig_img, resized_img, base_filename
-            )
+            base_filename = f"{pdf_stem}_page_{page_num}"
+            if self.crop_empty:
+                original_url, cropped_url, resized_url = self.cloudinary_uploader.upload_original_cropped_and_resized(
+                    raw_img, embed_img, resized_img, base_filename
+                )
+            else:
+                original_url, resized_url = self.cloudinary_uploader.upload_original_and_resized(
+                    raw_img, resized_img, base_filename
+                )
         
         # Sanitize text
         safe_text = self._sanitize_text(text[:10000]) if text else ""
         
-        # Build metadata (everything needed for saliency)
         metadata = {
-            # Document info
-            "filename": pdf_path.name,
+            "filename": filename,
             "page_number": page_num,
             "total_pages": total_pages,
             "has_text": bool(text and text.strip()),
@@ -378,11 +440,14 @@ class ProcessingPipeline:
             # Image URLs
             "page": resized_url or "",  # For display
             "original_url": original_url or "",
+            "cropped_url": cropped_url or "",
             "resized_url": resized_url or "",
             
             # Dimensions (needed for saliency overlay)
-            "original_width": orig_img.width,
-            "original_height": orig_img.height,
+            "original_width": raw_img.width,
+            "original_height": raw_img.height,
+            "cropped_width": int(embed_img.width) if self.crop_empty else int(raw_img.width),
+            "cropped_height": int(embed_img.height) if self.crop_empty else int(raw_img.height),
             "resized_width": resized_img.width,
             "resized_height": resized_img.height,
             
@@ -399,6 +464,15 @@ class ProcessingPipeline:
             
             # Strategy used (important for paper comparison)
             "embedding_strategy": self.embedding_strategy,
+
+            "model_name": getattr(self.embedder, "model_name", None),
+
+            "crop_empty_enabled": bool(self.crop_empty),
+            "crop_empty_crop_box": (crop_meta or {}).get("crop_box"),
+            "crop_empty_remove_page_number": bool(self.crop_empty_remove_page_number),
+            "crop_empty_percentage_to_remove": float(self.crop_empty_percentage_to_remove),
+            "crop_empty_preserve_border_px": int(self.crop_empty_preserve_border_px),
+            "crop_empty_uniform_rowcol_std_threshold": float(self.crop_empty_uniform_rowcol_std_threshold),
             
             # Extra metadata (year, district, etc.)
             **extra_metadata,
@@ -408,8 +482,10 @@ class ProcessingPipeline:
             "id": chunk_id,
             "visual_embedding": embedding_for_initial,    # "initial" vector in Qdrant
             "tile_pooled_embedding": embedding_for_pooling,  # "mean_pooling" vector in Qdrant
+            "experimental_pooled_embedding": experimental_pooled,  # "experimental_pooling" vector in Qdrant
+            "global_pooled_embedding": global_pooling,  # "global_pooling" vector in Qdrant
             "metadata": metadata,
-            "image": orig_img,
+            "image": raw_img,
             "resized_image": resized_img,
         }
         

visual_rag/indexing/qdrant_indexer.py

@@ -15,6 +15,7 @@ import time
 import hashlib
 import logging
 from typing import List, Dict, Any, Optional, Set
+from urllib.parse import urlparse
 import numpy as np
 
 logger = logging.getLogger(__name__)
@@ -54,6 +55,7 @@ class QdrantIndexer:
         collection_name: str,
         timeout: int = 60,
         prefer_grpc: bool = False,
+        vector_datatype: str = "float32",
     ):
         try:
             from qdrant_client import QdrantClient
@@ -65,17 +67,45 @@ class QdrantIndexer:
         
         self.collection_name = collection_name
         self.timeout = timeout
+        if vector_datatype not in ("float32", "float16"):
+            raise ValueError("vector_datatype must be 'float32' or 'float16'")
+        self.vector_datatype = vector_datatype
+        self._np_vector_dtype = np.float16 if vector_datatype == "float16" else np.float32
+
+        grpc_port = None
+        if prefer_grpc:
+            try:
+                parsed = urlparse(url)
+                port = parsed.port
+                if port == 6333:
+                    grpc_port = 6334
+            except Exception:
+                grpc_port = None
         
-        self.client = QdrantClient(
-            url=url,
-            api_key=api_key,
-            timeout=timeout,
-            prefer_grpc=prefer_grpc,
-            check_compatibility=False,
-        )
+        def _make_client(use_grpc: bool):
+            return QdrantClient(
+                url=url,
+                api_key=api_key,
+                timeout=timeout,
+                prefer_grpc=bool(use_grpc),
+                grpc_port=grpc_port,
+                check_compatibility=False,
+            )
+
+        self.client = _make_client(prefer_grpc)
+        if prefer_grpc:
+            try:
+                _ = self.client.get_collections()
+            except Exception as e:
+                msg = str(e)
+                if "StatusCode.PERMISSION_DENIED" in msg or "http2 header with status: 403" in msg:
+                    self.client = _make_client(False)
+                else:
+                    raise
         
         logger.info(f"🔌 Connected to Qdrant: {url}")
         logger.info(f"   Collection: {collection_name}")
+        logger.info(f"   Vector datatype: {self.vector_datatype}")
     
     def collection_exists(self) -> bool:
         """Check if collection exists."""
@@ -87,31 +117,29 @@ class QdrantIndexer:
         embedding_dim: int = 128,
         force_recreate: bool = False,
         enable_quantization: bool = False,
+        indexing_threshold: int = 20000,
+        full_scan_threshold: int = 0,
     ) -> bool:
         """
         Create collection with multi-vector support.
         
-        Creates two named vectors:
+        Creates named vectors:
         - initial: Full multi-vector embeddings (num_patches × dim)
         - mean_pooling: Tile-level pooled vectors (num_tiles × dim)
+        - experimental_pooling: Experimental multi-vector pooling (varies by model)
+        - global_pooling: Single vector pooled representation (dim)
         
         Args:
             embedding_dim: Embedding dimension (128 for ColSmol)
             force_recreate: Delete and recreate if exists
             enable_quantization: Enable int8 quantization
+            indexing_threshold: Qdrant optimizer indexing threshold (set 0 to always build ANN indexes)
         
         Returns:
             True if created, False if already existed
         """
         from qdrant_client.http import models
-        from qdrant_client.http.models import (
-            Distance,
-            VectorParams,
-            OptimizersConfigDiff,
-            HnswConfigDiff,
-            ScalarQuantizationConfig,
-            ScalarType,
-        )
+        from qdrant_client.http.models import Distance, VectorParams
         
         if self.collection_exists():
             if force_recreate:
@@ -128,54 +156,41 @@ class QdrantIndexer:
             comparator=models.MultiVectorComparator.MAX_SIM
         )
         
-        # HNSW config for pooled vectors
-        hnsw_config = HnswConfigDiff(
-            m=32,
-            ef_construct=100,
-            full_scan_threshold=10000,
-            on_disk=True,
-        )
-        
-        # Optional quantization
-        quantization_config = None
-        if enable_quantization:
-            logger.info("   Quantization: ENABLED (int8)")
-            quantization_config = ScalarQuantizationConfig(
-                type=ScalarType.INT8,
-                quantile=0.99,
-                always_ram=True,
-            )
-        
-        # Vector configs
+        # Vector configs - simplified for compatibility
+        datatype = models.Datatype.FLOAT16 if self.vector_datatype == "float16" else models.Datatype.FLOAT32
         vectors_config = {
             "initial": VectorParams(
                 size=embedding_dim,
                 distance=Distance.COSINE,
                 on_disk=True,
                 multivector_config=multivector_config,
-                quantization_config=quantization_config,
+                datatype=datatype,
             ),
             "mean_pooling": VectorParams(
                 size=embedding_dim,
                 distance=Distance.COSINE,
-                on_disk=False,  # Keep in RAM for fast prefetch
+                on_disk=False,
                 multivector_config=multivector_config,
-                hnsw_config=hnsw_config,
-                quantization_config=quantization_config,
+                datatype=datatype,
+            ),
+            "experimental_pooling": VectorParams(
+                size=embedding_dim,
+                distance=Distance.COSINE,
+                on_disk=False,
+                multivector_config=multivector_config,
+                datatype=datatype,
+            ),
+            "global_pooling": VectorParams(
+                size=embedding_dim,
+                distance=Distance.COSINE,
+                on_disk=False,
+                datatype=datatype,
             ),
         }
         
-        # Optimizer config for low-RAM clusters
-        optimizer_config = OptimizersConfigDiff(
-            indexing_threshold=20000,  # High threshold to avoid indexing (saves RAM)
-            memmap_threshold=0,  # Use mmap immediately
-            flush_interval_sec=5,  # Flush WAL frequently
-        )
-        
         self.client.create_collection(
             collection_name=self.collection_name,
             vectors_config=vectors_config,
-            optimizers_config=optimizer_config,
         )
         
         logger.info(f"✅ Collection created: {self.collection_name}")
@@ -202,16 +217,8 @@ class QdrantIndexer:
             "text": models.PayloadSchemaType.TEXT,
         }
         
-        # Default fields
-        if fields is None:
-            fields = [
-                {"field": "filename", "type": "keyword"},
-                {"field": "page_number", "type": "integer"},
-                {"field": "year", "type": "integer"},
-                {"field": "source", "type": "keyword"},
-                {"field": "district", "type": "keyword"},
-                {"field": "has_text", "type": "bool"},
-            ]
+        if not fields:
+            return
         
         logger.info("📇 Creating payload indexes...")
         
@@ -235,6 +242,8 @@ class QdrantIndexer:
         points: List[Dict[str, Any]],
         max_retries: int = 3,
         delay_between_batches: float = 0.5,
+        wait: bool = True,
+        stop_event=None,
     ) -> int:
         """
         Upload a batch of points to Qdrant.
@@ -243,12 +252,16 @@ class QdrantIndexer:
         - id: Unique point ID (string or UUID)
         - visual_embedding: Full embedding [num_patches, dim]
         - tile_pooled_embedding: Pooled embedding [num_tiles, dim]
+        - experimental_pooled_embedding: Experimental pooled embedding [*, dim]
+        - global_pooled_embedding: Pooled embedding [dim]
         - metadata: Payload dict
         
         Args:
             points: List of point dicts
             max_retries: Retry attempts on failure
             delay_between_batches: Delay after upload
+            wait: Wait for operation to complete on Qdrant server
+            stop_event: Optional threading.Event used to cancel uploads early
         
         Returns:
             Number of successfully uploaded points
@@ -257,39 +270,115 @@ class QdrantIndexer:
         
         if not points:
             return 0
+
+        def _is_cancelled() -> bool:
+            return stop_event is not None and getattr(stop_event, "is_set", lambda: False)()
         
-        # Build Qdrant points
-        qdrant_points = []
-        for point_data in points:
-            point = models.PointStruct(
-                id=point_data["id"],
-                vector={
-                    "initial": point_data["visual_embedding"].tolist(),
-                    "mean_pooling": point_data["tile_pooled_embedding"].tolist(),
-                },
-                payload=point_data["metadata"],
-            )
-            qdrant_points.append(point)
+        def _is_payload_too_large_error(e: Exception) -> bool:
+            msg = str(e)
+            if ("JSON payload" in msg and "larger than allowed" in msg) or ("Payload error:" in msg and "limit:" in msg):
+                return True
+            content = getattr(e, "content", None)
+            if content is not None:
+                try:
+                    if isinstance(content, (bytes, bytearray)):
+                        text = content.decode("utf-8", errors="ignore")
+                    else:
+                        text = str(content)
+                except Exception:
+                    text = ""
+                if ("JSON payload" in text and "larger than allowed" in text) or ("Payload error" in text and "limit" in text):
+                    return True
+            resp = getattr(e, "response", None)
+            if resp is not None:
+                try:
+                    text = str(getattr(resp, "text", "") or "")
+                except Exception:
+                    text = ""
+                if ("JSON payload" in text and "larger than allowed" in text) or ("Payload error" in text and "limit" in text):
+                    return True
+            return False
+
+        def _to_list(val):
+            if isinstance(val, np.ndarray):
+                return val.tolist()
+            return val
+
+        def _build_qdrant_points(batch_points: List[Dict[str, Any]]) -> List[models.PointStruct]:
+            qdrant_points: List[models.PointStruct] = []
+            for p in batch_points:
+                global_pooled = p.get("global_pooled_embedding")
+                if global_pooled is None:
+                    tile_pooled = np.array(p["tile_pooled_embedding"], dtype=np.float32)
+                    global_pooled = tile_pooled.mean(axis=0)
+                global_pooled = np.array(global_pooled, dtype=np.float32).reshape(-1)
+
+                initial = np.array(p["visual_embedding"], dtype=np.float32).astype(self._np_vector_dtype, copy=False)
+                mean_pooling = np.array(p["tile_pooled_embedding"], dtype=np.float32).astype(self._np_vector_dtype, copy=False)
+                experimental_pooling = np.array(p["experimental_pooled_embedding"], dtype=np.float32).astype(
+                    self._np_vector_dtype, copy=False
+                )
+                global_pooling = global_pooled.astype(self._np_vector_dtype, copy=False)
+
+                qdrant_points.append(
+                    models.PointStruct(
+                        id=p["id"],
+                        vector={
+                            "initial": _to_list(initial),
+                            "mean_pooling": _to_list(mean_pooling),
+                            "experimental_pooling": _to_list(experimental_pooling),
+                            "global_pooling": _to_list(global_pooling),
+                        },
+                        payload=p["metadata"],
+                    )
+                )
+            return qdrant_points
         
         # Upload with retry
         for attempt in range(max_retries):
             try:
+                if _is_cancelled():
+                    return 0
+                qdrant_points = _build_qdrant_points(points)
                 self.client.upsert(
                     collection_name=self.collection_name,
                     points=qdrant_points,
-                    wait=True,
+                    wait=wait,
                 )
-                
+
                 if delay_between_batches > 0:
+                    if _is_cancelled():
+                        return 0
                     time.sleep(delay_between_batches)
-                
-                return len(qdrant_points)
-                
+
+                return len(points)
+
             except Exception as e:
-                logger.warning(
-                    f"Upload attempt {attempt + 1}/{max_retries} failed: {e}"
-                )
+                if _is_payload_too_large_error(e) and len(points) > 1:
+                    mid = len(points) // 2
+                    left = points[:mid]
+                    right = points[mid:]
+                    logger.warning(
+                        f"Upload payload too large for {len(points)} points; splitting into {len(left)} + {len(right)}"
+                    )
+                    return self.upload_batch(
+                        left,
+                        max_retries=max_retries,
+                        delay_between_batches=delay_between_batches,
+                        wait=wait,
+                        stop_event=stop_event,
+                    ) + self.upload_batch(
+                        right,
+                        max_retries=max_retries,
+                        delay_between_batches=delay_between_batches,
+                        wait=wait,
+                        stop_event=stop_event,
+                    )
+
+                logger.warning(f"Upload attempt {attempt + 1}/{max_retries} failed: {e}")
                 if attempt < max_retries - 1:
+                    if _is_cancelled():
+                        return 0
                     time.sleep(2 ** attempt)  # Exponential backoff
         
         logger.error(f"❌ Upload failed after {max_retries} attempts")

visual_rag/preprocessing/__init__.py

@@ -0,0 +1,5 @@
+from visual_rag.preprocessing.crop_empty import CropEmptyConfig, crop_empty
+
+__all__ = ["CropEmptyConfig", "crop_empty"]
+
+

visual_rag/preprocessing/crop_empty.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Tuple
+
+import numpy as np
+from PIL import Image
+
+
+@dataclass(frozen=True)
+class CropEmptyConfig:
+    percentage_to_remove: float = 0.9
+    remove_page_number: bool = False
+    color_threshold: int = 240
+    min_white_fraction: float = 0.99
+    content_density_sides: float = 0.001
+    content_density_main_text: float = 0.05
+    content_density_any: float = 1e-6
+    preserve_border_px: int = 1
+    uniform_rowcol_std_threshold: float = 0.0
+
+
+def crop_empty(image: Image.Image, *, config: CropEmptyConfig) -> Tuple[Image.Image, Dict[str, Any]]:
+    img = image.convert("RGB")
+    arr = np.array(img)
+    intensity = arr.mean(axis=2)
+
+    def _find_border_start(axis: int, *, min_content_density_threshold: float) -> int:
+        size = intensity.shape[axis]
+        for i in range(size):
+            pixels = intensity[i, :] if axis == 0 else intensity[:, i]
+            white = float(np.mean(pixels > config.color_threshold))
+            non_white = 1.0 - white
+            if float(config.uniform_rowcol_std_threshold) > 0.0 and float(np.std(pixels)) <= float(config.uniform_rowcol_std_threshold):
+                continue
+            if (white < config.min_white_fraction) and (non_white > min_content_density_threshold):
+                return int(i)
+        return int(size)
+
+    def _find_border_end(axis: int, *, min_content_density_threshold: float) -> int:
+        size = intensity.shape[axis]
+        for i in range(size - 1, -1, -1):
+            pixels = intensity[i, :] if axis == 0 else intensity[:, i]
+            white = float(np.mean(pixels > config.color_threshold))
+            non_white = 1.0 - white
+            if float(config.uniform_rowcol_std_threshold) > 0.0 and float(np.std(pixels)) <= float(config.uniform_rowcol_std_threshold):
+                continue
+            if (white < config.min_white_fraction) and (non_white > min_content_density_threshold):
+                return int(i + 1)
+        return 0
+
+    top = _find_border_start(0, min_content_density_threshold=float(config.content_density_sides))
+    left = _find_border_start(1, min_content_density_threshold=float(config.content_density_sides))
+    right = _find_border_end(1, min_content_density_threshold=float(config.content_density_sides))
+
+    main_text_end = _find_border_end(0, min_content_density_threshold=float(config.content_density_main_text))
+    last_content_end = _find_border_end(0, min_content_density_threshold=float(config.content_density_any))
+    bottom = main_text_end if config.remove_page_number else last_content_end
+
+    width, height = img.size
+    pad = max(int(getattr(config, "preserve_border_px", 0) or 0), 0)
+    if pad > 0:
+        left = max(int(left) - pad, 0)
+        top = max(int(top) - pad, 0)
+        right = min(int(right) + pad, int(width))
+        bottom = min(int(bottom) + pad, int(height))
+    crop_box = (int(left), int(top), int(right), int(bottom))
+    valid = 0 <= crop_box[0] < crop_box[2] <= width and 0 <= crop_box[1] < crop_box[3] <= height
+
+    if not valid:
+        return image, {
+            "applied": False,
+            "crop_box": None,
+            "original_width": int(width),
+            "original_height": int(height),
+            "cropped_width": int(width),
+            "cropped_height": int(height),
+            "config": {
+                "percentage_to_remove": float(config.percentage_to_remove),
+                "remove_page_number": bool(config.remove_page_number),
+                "color_threshold": int(config.color_threshold),
+                "min_white_fraction": float(config.min_white_fraction),
+                "content_density_sides": float(config.content_density_sides),
+                "content_density_main_text": float(config.content_density_main_text),
+                "content_density_any": float(config.content_density_any),
+                "preserve_border_px": int(config.preserve_border_px),
+                "uniform_rowcol_std_threshold": float(config.uniform_rowcol_std_threshold),
+            },
+        }
+
+    cropped = img.crop(crop_box)
+    return cropped, {
+        "applied": True,
+        "crop_box": [int(crop_box[0]), int(crop_box[1]), int(crop_box[2]), int(crop_box[3])],
+        "original_width": int(width),
+        "original_height": int(height),
+        "cropped_width": int(cropped.width),
+        "cropped_height": int(cropped.height),
+        "config": {
+            "percentage_to_remove": float(config.percentage_to_remove),
+            "remove_page_number": bool(config.remove_page_number),
+            "color_threshold": int(config.color_threshold),
+            "min_white_fraction": float(config.min_white_fraction),
+            "content_density_sides": float(config.content_density_sides),
+            "content_density_main_text": float(config.content_density_main_text),
+            "content_density_any": float(config.content_density_any),
+            "preserve_border_px": int(config.preserve_border_px),
+            "uniform_rowcol_std_threshold": float(config.uniform_rowcol_std_threshold),
+        },
+    }
+
+

visual_rag/qdrant_admin.py

@@ -0,0 +1,205 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, Optional
+from urllib.parse import urlparse
+
+
+@dataclass(frozen=True)
+class QdrantConnection:
+    url: str
+    api_key: Optional[str]
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except Exception:
+        return
+    try:
+        from pathlib import Path
+
+        if Path(".env").exists():
+            load_dotenv(".env")
+    except Exception:
+        return
+
+
+def _resolve_qdrant_connection(
+    *,
+    url: Optional[str] = None,
+    api_key: Optional[str] = None,
+) -> QdrantConnection:
+    import os
+
+    _maybe_load_dotenv()
+    resolved_url = url or os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    if not resolved_url:
+        raise ValueError("Qdrant URL not set (pass url= or set SIGIR_QDRANT_URL/DEST_QDRANT_URL/QDRANT_URL).")
+    resolved_key = (
+        api_key
+        or os.getenv("SIGIR_QDRANT_KEY")
+        or os.getenv("SIGIR_QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("QDRANT_API_KEY")
+    )
+    return QdrantConnection(url=str(resolved_url), api_key=resolved_key)
+
+
+def _infer_grpc_port(url: str) -> Optional[int]:
+    try:
+        if urlparse(url).port == 6333:
+            return 6334
+    except Exception:
+        return None
+    return None
+
+
+class QdrantAdmin:
+    def __init__(
+        self,
+        *,
+        url: Optional[str] = None,
+        api_key: Optional[str] = None,
+        prefer_grpc: bool = False,
+        timeout: int = 60,
+    ):
+        from qdrant_client import QdrantClient
+
+        conn = _resolve_qdrant_connection(url=url, api_key=api_key)
+        grpc_port = _infer_grpc_port(conn.url) if prefer_grpc else None
+        self.client = QdrantClient(
+            url=conn.url,
+            api_key=conn.api_key,
+            prefer_grpc=bool(prefer_grpc),
+            grpc_port=grpc_port,
+            timeout=int(timeout),
+            check_compatibility=False,
+        )
+
+    def get_collection_info(self, *, collection_name: str) -> Dict[str, Any]:
+        info = self.client.get_collection(collection_name)
+        try:
+            return info.model_dump()
+        except Exception:
+            try:
+                return info.dict()
+            except Exception:
+                return {"collection": str(collection_name), "raw": str(info)}
+
+    def modify_collection_config(
+        self,
+        *,
+        collection_name: str,
+        hnsw_config: Optional[Dict[str, Any]] = None,
+        collection_params: Optional[Dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+    ) -> bool:
+        """
+        Patch collection-level config via Qdrant update_collection.
+
+        Supported keys:
+        - hnsw_config: dict for HnswConfigDiff (e.g. on_disk, m, ef_construct, full_scan_threshold)
+        - collection_params: dict for CollectionParamsDiff (e.g. on_disk_payload)
+        """
+        from qdrant_client.http import models as m
+
+        hnsw_diff = m.HnswConfigDiff(**hnsw_config) if isinstance(hnsw_config, dict) else None
+        params_diff = m.CollectionParamsDiff(**collection_params) if isinstance(collection_params, dict) else None
+        if hnsw_diff is None and params_diff is None:
+            raise ValueError("No changes provided (pass hnsw_config and/or collection_params).")
+        return bool(
+            self.client.update_collection(
+                collection_name=str(collection_name),
+                hnsw_config=hnsw_diff,
+                collection_params=params_diff,
+                timeout=int(timeout) if timeout is not None else None,
+            )
+        )
+
+    def modify_collection_vector_config(
+        self,
+        *,
+        collection_name: str,
+        vectors: Dict[str, Dict[str, Any]],
+        timeout: Optional[int] = None,
+    ) -> bool:
+        """
+        Patch vector params under params.vectors[vector_name] using Qdrant update_collection.
+
+        Supported keys per vector:
+        - on_disk: bool
+        - hnsw_config: dict with optional keys: m, ef_construct, full_scan_threshold, on_disk
+        """
+        from qdrant_client.http import models as m
+
+        collection_name = str(collection_name)
+        info = self.client.get_collection(collection_name)
+        existing = set()
+        try:
+            existing = set((info.config.params.vectors or {}).keys())
+        except Exception:
+            existing = set()
+
+        missing = [str(k) for k in (vectors or {}).keys() if existing and str(k) not in existing]
+        if missing:
+            raise ValueError(f"Vectors do not exist in collection '{collection_name}': {missing}. Existing: {sorted(existing)}")
+
+        ok = True
+        for name, cfg in (vectors or {}).items():
+            if not isinstance(cfg, dict):
+                raise ValueError(f"vectors['{name}'] must be a dict, got {type(cfg)}")
+            hnsw_cfg = cfg.get("hnsw_config")
+            hnsw_diff = m.HnswConfigDiff(**hnsw_cfg) if isinstance(hnsw_cfg, dict) else None
+            vectors_diff = {
+                str(name): m.VectorParamsDiff(
+                    on_disk=cfg.get("on_disk", None),
+                    hnsw_config=hnsw_diff,
+                )
+            }
+
+            ok = bool(
+                self.client.update_collection(
+                    collection_name=collection_name,
+                    vectors_config=vectors_diff,
+                    timeout=int(timeout) if timeout is not None else None,
+                )
+            ) and ok
+
+        return ok
+
+    def ensure_collection_all_on_disk(
+        self,
+        *,
+        collection_name: str,
+        timeout: Optional[int] = None,
+    ) -> Dict[str, Any]:
+        """
+        Ensure:
+        - All existing named vectors have on_disk=True and hnsw_config.on_disk=True
+        - Collection hnsw_config.on_disk=True
+        - Collection params.on_disk_payload=True
+        Returns the post-update collection info (dict).
+        """
+        collection_name = str(collection_name)
+        info = self.client.get_collection(collection_name)
+        vectors = {}
+        try:
+            existing = list((info.config.params.vectors or {}).keys())
+        except Exception:
+            existing = []
+        for vname in existing:
+            vectors[str(vname)] = {"on_disk": True, "hnsw_config": {"on_disk": True}}
+
+        if vectors:
+            self.modify_collection_vector_config(collection_name=collection_name, vectors=vectors, timeout=timeout)
+
+        self.modify_collection_config(
+            collection_name=collection_name,
+            hnsw_config={"on_disk": True},
+            collection_params={"on_disk_payload": True},
+            timeout=timeout,
+        )
+
+        return self.get_collection_info(collection_name=collection_name)
+

visual_rag/retrieval/__init__.py

@@ -7,7 +7,13 @@ Components:
 """
 
 from visual_rag.retrieval.two_stage import TwoStageRetriever
+from visual_rag.retrieval.single_stage import SingleStageRetriever
+from visual_rag.retrieval.multi_vector import MultiVectorRetriever
+from visual_rag.retrieval.three_stage import ThreeStageRetriever
 
 __all__ = [
     "TwoStageRetriever",
+    "SingleStageRetriever",
+    "MultiVectorRetriever",
+    "ThreeStageRetriever",
 ]

visual_rag/retrieval/multi_vector.py

@@ -0,0 +1,220 @@
+import os
+from typing import Any, Dict, List, Optional
+from urllib.parse import urlparse
+
+from visual_rag.embedding.visual_embedder import VisualEmbedder
+from visual_rag.retrieval.single_stage import SingleStageRetriever
+from visual_rag.retrieval.two_stage import TwoStageRetriever
+from visual_rag.retrieval.three_stage import ThreeStageRetriever
+
+
+class MultiVectorRetriever:
+    @staticmethod
+    def _maybe_load_dotenv() -> None:
+        try:
+            from dotenv import load_dotenv
+        except ImportError:
+            return
+        if os.path.exists(".env"):
+            load_dotenv(".env")
+
+    def __init__(
+        self,
+        collection_name: str,
+        model_name: str = "vidore/colSmol-500M",
+        qdrant_url: Optional[str] = None,
+        qdrant_api_key: Optional[str] = None,
+        prefer_grpc: bool = False,
+        request_timeout: int = 120,
+        max_retries: int = 3,
+        retry_sleep: float = 0.5,
+        qdrant_client=None,
+        embedder: Optional[VisualEmbedder] = None,
+    ):
+        if qdrant_client is None:
+            self._maybe_load_dotenv()
+            try:
+                from qdrant_client import QdrantClient
+            except ImportError as e:
+                raise ImportError(
+                    "Qdrant client not installed. Install with: pip install visual-rag-toolkit[qdrant]"
+                ) from e
+
+            qdrant_url = (
+                qdrant_url
+                or os.getenv("SIGIR_QDRANT_URL")
+                or os.getenv("DEST_QDRANT_URL")
+                or os.getenv("QDRANT_URL")
+            )
+            if not qdrant_url:
+                raise ValueError(
+                    "QDRANT_URL is required (pass qdrant_url or set env var). "
+                    "You can also set DEST_QDRANT_URL to override."
+                )
+
+            qdrant_api_key = (
+                qdrant_api_key
+                or os.getenv("SIGIR_QDRANT_KEY")
+                or os.getenv("SIGIR_QDRANT_API_KEY")
+                or os.getenv("DEST_QDRANT_API_KEY")
+                or os.getenv("QDRANT_API_KEY")
+            )
+
+            grpc_port = None
+            if prefer_grpc:
+                try:
+                    if urlparse(qdrant_url).port == 6333:
+                        grpc_port = 6334
+                except Exception:
+                    grpc_port = None
+            def _make_client(use_grpc: bool):
+                return QdrantClient(
+                    url=qdrant_url,
+                    api_key=qdrant_api_key,
+                    prefer_grpc=bool(use_grpc),
+                    grpc_port=grpc_port,
+                    timeout=int(request_timeout),
+                    check_compatibility=False,
+                )
+
+            qdrant_client = _make_client(prefer_grpc)
+            if prefer_grpc:
+                try:
+                    _ = qdrant_client.get_collections()
+                except Exception as e:
+                    msg = str(e)
+                    if "StatusCode.PERMISSION_DENIED" in msg or "http2 header with status: 403" in msg:
+                        qdrant_client = _make_client(False)
+                    else:
+                        raise
+
+        self.client = qdrant_client
+        self.collection_name = collection_name
+        self.embedder = embedder or VisualEmbedder(model_name=model_name)
+
+        self._two_stage = TwoStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+            max_retries=int(max_retries),
+            retry_sleep=float(retry_sleep),
+        )
+        self._three_stage = ThreeStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+            max_retries=int(max_retries),
+            retry_sleep=float(retry_sleep),
+        )
+        self._single_stage = SingleStageRetriever(
+            self.client,
+            collection_name=self.collection_name,
+            request_timeout=int(request_timeout),
+        )
+
+    def build_filter(
+        self,
+        year: Optional[Any] = None,
+        source: Optional[str] = None,
+        district: Optional[str] = None,
+        filename: Optional[str] = None,
+        has_text: Optional[bool] = None,
+    ):
+        return self._two_stage.build_filter(
+            year=year,
+            source=source,
+            district=district,
+            filename=filename,
+            has_text=has_text,
+        )
+
+    def search(
+        self,
+        query: str,
+        top_k: int = 10,
+        mode: str = "single_full",
+        prefetch_k: Optional[int] = None,
+        stage1_mode: str = "pooled_query_vs_tiles",
+        filter_obj=None,
+        return_embeddings: bool = False,
+    ) -> List[Dict[str, Any]]:
+        q = self.embedder.embed_query(query)
+        try:
+            import torch
+        except ImportError:
+            torch = None
+        if torch is not None and isinstance(q, torch.Tensor):
+            query_embedding = q.detach().cpu().numpy()
+        else:
+            query_embedding = q.numpy()
+
+        return self.search_embedded(
+            query_embedding=query_embedding,
+            top_k=top_k,
+            mode=mode,
+            prefetch_k=prefetch_k,
+            stage1_mode=stage1_mode,
+            filter_obj=filter_obj,
+            return_embeddings=return_embeddings,
+        )
+
+    def search_embedded(
+        self,
+        *,
+        query_embedding,
+        top_k: int = 10,
+        mode: str = "single_full",
+        prefetch_k: Optional[int] = None,
+        stage1_mode: str = "pooled_query_vs_tiles",
+        stage1_k: Optional[int] = None,
+        stage2_k: Optional[int] = None,
+        filter_obj=None,
+        return_embeddings: bool = False,
+    ) -> List[Dict[str, Any]]:
+        if mode == "single_full":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="multi_vector",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "single_tiles":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="tiles_maxsim",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "single_global":
+            return self._single_stage.search(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                strategy="pooled_global",
+                filter_obj=filter_obj,
+            )
+
+        if mode == "two_stage":
+            return self._two_stage.search_server_side(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                prefetch_k=prefetch_k,
+                filter_obj=filter_obj,
+                stage1_mode=stage1_mode,
+            )
+
+        if mode == "three_stage":
+            s1 = int(stage1_k) if stage1_k is not None else 1000
+            s2 = int(stage2_k) if stage2_k is not None else 300
+            return self._three_stage.search_server_side(
+                query_embedding=query_embedding,
+                top_k=top_k,
+                stage1_k=s1,
+                stage2_k=s2,
+                filter_obj=filter_obj,
+            )
+
+        raise ValueError(f"Unknown mode: {mode}")
+
+

visual_rag/retrieval/single_stage.py

@@ -20,10 +20,11 @@ class SingleStageRetriever:
     """
     Single-stage visual document retrieval using native Qdrant search.
     
-    Supports three strategies:
-    1. multi_vector (SOTA): Native MaxSim on full embeddings
-    2. pooled_tile: Search on tile-level pooled vectors
-    3. pooled_global: Search on globally pooled vectors
+    Supports strategies:
+    - multi_vector: Native MaxSim on full embeddings (using="initial")
+    - tiles_maxsim: Native MaxSim between query tokens and tile vectors (using="mean_pooling")
+    - pooled_tile: Pooled query vs tile vectors (using="mean_pooling")
+    - pooled_global: Pooled query vs global pooled doc vector (using="global_pooling")
     
     Args:
         qdrant_client: Connected Qdrant client
@@ -38,9 +39,11 @@ class SingleStageRetriever:
         self,
         qdrant_client,
         collection_name: str,
+        request_timeout: int = 120,
     ):
         self.client = qdrant_client
         self.collection_name = collection_name
+        self.request_timeout = int(request_timeout)
     
     def search(
         self,
@@ -55,7 +58,7 @@ class SingleStageRetriever:
         Args:
             query_embedding: Query embeddings [num_tokens, dim]
             top_k: Number of results
-            strategy: "multi_vector", "pooled_tile", or "pooled_global"
+            strategy: "multi_vector", "tiles_maxsim", "pooled_tile", or "pooled_global"
             filter_obj: Qdrant filter
         
         Returns:
@@ -67,21 +70,27 @@ class SingleStageRetriever:
             # Native multi-vector MaxSim
             vector_name = "initial"
             query_vector = query_np.tolist()
-            logger.info(f"🎯 Multi-vector search on '{vector_name}'")
+            logger.debug(f"🎯 Multi-vector search on '{vector_name}'")
+            
+        elif strategy == "tiles_maxsim":
+            # Native multi-vector MaxSim against tile vectors
+            vector_name = "mean_pooling"
+            query_vector = query_np.tolist()
+            logger.debug(f"🎯 Tile MaxSim search on '{vector_name}'")
             
         elif strategy == "pooled_tile":
             # Tile-level pooled
             vector_name = "mean_pooling"
             query_pooled = query_np.mean(axis=0)
             query_vector = query_pooled.tolist()
-            logger.info(f"🔍 Tile-pooled search on '{vector_name}'")
+            logger.debug(f"🔍 Tile-pooled search on '{vector_name}'")
             
         elif strategy == "pooled_global":
-            # Global mean pooling (single vector)
-            vector_name = "mean_pooling"
+            # Global pooled vector (single vector)
+            vector_name = "global_pooling"
             query_pooled = query_np.mean(axis=0)
             query_vector = query_pooled.tolist()
-            logger.info(f"🔍 Global-pooled search on '{vector_name}'")
+            logger.debug(f"🔍 Global-pooled search on '{vector_name}'")
             
         else:
             raise ValueError(f"Unknown strategy: {strategy}")
@@ -94,7 +103,7 @@ class SingleStageRetriever:
             limit=top_k,
             with_payload=True,
             with_vectors=False,
-            timeout=120,
+            timeout=self.request_timeout,
         ).points
         
         return [

visual_rag/retrieval/three_stage.py

@@ -0,0 +1,174 @@
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class ThreeStageRetriever:
+    def __init__(
+        self,
+        qdrant_client,
+        collection_name: str,
+        *,
+        full_vector_name: str = "initial",
+        experimental_vector_name: str = "experimental_pooling",
+        global_vector_name: str = "global_pooling",
+        request_timeout: int = 120,
+        max_retries: int = 3,
+        retry_sleep: float = 0.5,
+    ):
+        self.client = qdrant_client
+        self.collection_name = collection_name
+        self.full_vector_name = full_vector_name
+        self.experimental_vector_name = experimental_vector_name
+        self.global_vector_name = global_vector_name
+        self.request_timeout = int(request_timeout)
+        self.max_retries = int(max_retries)
+        self.retry_sleep = float(retry_sleep)
+
+        self._global_is_multivector: Optional[bool] = None
+        self._experimental_is_multivector: Optional[bool] = None
+
+    def _retry_call(self, fn):
+        import time
+
+        last_err = None
+        for attempt in range(self.max_retries):
+            try:
+                return fn()
+            except Exception as e:
+                last_err = e
+                if attempt >= self.max_retries - 1:
+                    break
+                time.sleep(self.retry_sleep * (2 ** attempt))
+        if last_err is not None:
+            raise last_err
+
+    def _to_numpy(self, embedding: Union[torch.Tensor, np.ndarray]) -> np.ndarray:
+        if isinstance(embedding, torch.Tensor):
+            if embedding.dtype == torch.bfloat16:
+                return embedding.cpu().float().numpy()
+            return embedding.cpu().numpy()
+        return np.array(embedding, dtype=np.float32)
+
+    def _infer_vector_is_multivector(self, vector_name: str) -> bool:
+        info = self.client.get_collection(self.collection_name)
+        cfg = getattr(info, "config", None)
+        params = getattr(cfg, "params", None) if cfg is not None else None
+        vectors = getattr(params, "vectors", None) if params is not None else None
+        v = None
+        try:
+            if isinstance(vectors, dict):
+                v = vectors.get(vector_name)
+            else:
+                v = vectors[vector_name]
+        except Exception:
+            v = None
+        mv = getattr(v, "multivector_config", None) if v is not None else None
+        if mv is None and isinstance(v, dict):
+            mv = v.get("multivector_config")
+        return mv is not None
+
+    def _and_filter(self, base_filter, ids: List[Any]):
+        from qdrant_client.http import models as m
+
+        has_id = m.HasIdCondition(has_id=list(ids))
+        if base_filter is None:
+            return m.Filter(must=[has_id])
+        return m.Filter(must=[base_filter, has_id])
+
+    def search_server_side(
+        self,
+        *,
+        query_embedding: Union[torch.Tensor, np.ndarray],
+        top_k: int = 100,
+        stage1_k: int = 1000,
+        stage2_k: int = 300,
+        filter_obj=None,
+    ) -> List[Dict[str, Any]]:
+        from qdrant_client.http import models as m
+
+        query_np = self._to_numpy(query_embedding)
+
+        stage1_query = query_np.mean(axis=0).tolist()
+        stage2_query = query_np.tolist()
+        stage3_query = query_np.tolist()
+
+        logger.info(f"Stage 1: global prefetch {int(stage1_k)}")
+
+        def _do_stage1():
+            return self.client.query_points(
+                collection_name=self.collection_name,
+                query=stage1_query,
+                using=self.global_vector_name,
+                limit=int(stage1_k),
+                query_filter=filter_obj,
+                with_payload=False,
+                with_vectors=False,
+                timeout=self.request_timeout,
+            ).points
+
+        s1 = self._retry_call(_do_stage1)
+        if not s1:
+            return []
+        s1_ids = [p.id for p in s1]
+        s1_score = {str(p.id): float(p.score) for p in s1}
+
+        logger.info(f"Stage 2: experimental prefetch {int(stage2_k)} (restricted to stage1)")
+
+        stage2_filter = self._and_filter(filter_obj, s1_ids)
+
+        def _do_stage2():
+            return self.client.query_points(
+                collection_name=self.collection_name,
+                query=stage2_query,
+                using=self.experimental_vector_name,
+                limit=int(min(int(stage2_k), len(s1_ids))),
+                query_filter=stage2_filter,
+                with_payload=False,
+                with_vectors=False,
+                timeout=self.request_timeout,
+            ).points
+
+        s2 = self._retry_call(_do_stage2)
+        if not s2:
+            return []
+        s2_ids = [p.id for p in s2]
+        s2_score = {str(p.id): float(p.score) for p in s2}
+
+        logger.info(f"Stage 3: exact rerank on initial to top {int(top_k)} (restricted to stage2)")
+
+        stage3_filter = self._and_filter(filter_obj, s2_ids)
+
+        def _do_stage3():
+            return self.client.query_points(
+                collection_name=self.collection_name,
+                query=stage3_query,
+                using=self.full_vector_name,
+                limit=int(top_k),
+                query_filter=stage3_filter,
+                with_payload=True,
+                with_vectors=False,
+                search_params=m.SearchParams(exact=True),
+                timeout=self.request_timeout,
+            ).points
+
+        s3 = self._retry_call(_do_stage3)
+        out = []
+        for p in s3:
+            pid = str(p.id)
+            out.append(
+                {
+                    "id": p.id,
+                    "score_stage1": s1_score.get(pid),
+                    "score_stage2": s2_score.get(pid),
+                    "score_stage3": float(p.score),
+                    "score_final": float(p.score),
+                    "payload": p.payload,
+                }
+            )
+        return out
+

visual_rag/retrieval/two_stage.py

@@ -64,11 +64,119 @@ class TwoStageRetriever:
         collection_name: str,
         full_vector_name: str = "initial",
         pooled_vector_name: str = "mean_pooling",
+        experimental_vector_name: str = "experimental_pooling",
+        global_vector_name: str = "global_pooling",
+        request_timeout: int = 120,
+        max_retries: int = 3,
+        retry_sleep: float = 0.5,
     ):
         self.client = qdrant_client
         self.collection_name = collection_name
         self.full_vector_name = full_vector_name
         self.pooled_vector_name = pooled_vector_name
+        self.experimental_vector_name = experimental_vector_name
+        self.global_vector_name = global_vector_name
+        self.request_timeout = int(request_timeout)
+        self.max_retries = int(max_retries)
+        self.retry_sleep = float(retry_sleep)
+
+    def _retry_call(self, fn):
+        import time
+
+        last_err = None
+        for attempt in range(self.max_retries):
+            try:
+                return fn()
+            except Exception as e:
+                last_err = e
+                if attempt >= self.max_retries - 1:
+                    break
+                time.sleep(self.retry_sleep * (2 ** attempt))
+        if last_err is not None:
+            raise last_err
+
+    def search_server_side(
+        self,
+        query_embedding: Union[torch.Tensor, np.ndarray],
+        top_k: int = 10,
+        prefetch_k: Optional[int] = None,
+        filter_obj=None,
+        stage1_mode: str = "pooled_query_vs_tiles",
+    ) -> List[Dict[str, Any]]:
+        """
+        Two-stage retrieval using Qdrant's native prefetch (all server-side).
+        
+        This is MUCH faster than search() because it avoids network transfer
+        of large multi-vector embeddings. All computation happens in Qdrant.
+        
+        Args:
+            query_embedding: Query embeddings [num_tokens, dim]
+            top_k: Final number of results
+            prefetch_k: Candidates for stage 1 (default: 10x top_k)
+            filter_obj: Qdrant filter
+            stage1_mode: How to do stage 1 prefetch
+        
+        Returns:
+            List of results with scores
+        """
+        from qdrant_client.http import models
+        
+        query_np = self._to_numpy(query_embedding)
+        
+        if prefetch_k is None:
+            prefetch_k = max(100, top_k * 10)
+        
+        if stage1_mode == "pooled_query_vs_tiles":
+            prefetch_query = query_np.mean(axis=0).tolist()
+            prefetch_using = self.pooled_vector_name
+        elif stage1_mode == "tokens_vs_tiles":
+            prefetch_query = query_np.tolist()
+            prefetch_using = self.pooled_vector_name
+        elif stage1_mode == "pooled_query_vs_experimental":
+            prefetch_query = query_np.mean(axis=0).tolist()
+            prefetch_using = self.experimental_vector_name
+        elif stage1_mode == "tokens_vs_experimental":
+            prefetch_query = query_np.tolist()
+            prefetch_using = self.experimental_vector_name
+        elif stage1_mode == "pooled_query_vs_global":
+            prefetch_query = query_np.mean(axis=0).tolist()
+            prefetch_using = self.global_vector_name
+        else:
+            raise ValueError(f"Unknown stage1_mode: {stage1_mode}")
+        
+        rerank_query = query_np.tolist()
+        
+        def _do_query():
+            return self.client.query_points(
+                collection_name=self.collection_name,
+                query=rerank_query,
+                using=self.full_vector_name,
+                limit=top_k,
+                query_filter=filter_obj,
+                with_payload=True,
+                search_params=models.SearchParams(exact=True),
+                prefetch=[
+                    models.Prefetch(
+                        query=prefetch_query,
+                        using=prefetch_using,
+                        limit=prefetch_k,
+                    )
+                ],
+                timeout=self.request_timeout,
+            ).points
+        
+        results = self._retry_call(_do_query)
+        
+        return [
+            {
+                "id": r.id,
+                "score_stage1": None,
+                "score_stage2": r.score,
+                "score_final": r.score,
+                "payload": r.payload,
+            }
+            for r in results
+        ]
     
     def search(
         self,
@@ -78,6 +186,7 @@ class TwoStageRetriever:
         filter_obj=None,
         use_reranking: bool = True,
         return_embeddings: bool = False,
+        stage1_mode: str = "pooled_query_vs_tiles",
     ) -> List[Dict[str, Any]]:
         """
         Two-stage retrieval: prefetch with pooling, rerank with MaxSim.
@@ -89,6 +198,10 @@ class TwoStageRetriever:
             filter_obj: Qdrant filter for metadata filtering
             use_reranking: Enable stage 2 reranking (default: True)
             return_embeddings: Include embeddings in results
+            stage1_mode:
+                - "pooled_query_vs_tiles": pool query to 1×dim and search tile vectors (using="mean_pooling")
+                - "tokens_vs_tiles": search tile vectors with full query tokens (using="mean_pooling")
+                - "pooled_query_vs_global": pool query to 1×dim and search global pooled doc vectors (using="global_pooling")
         
         Returns:
             List of results with scores and metadata:
@@ -111,11 +224,12 @@ class TwoStageRetriever:
             prefetch_k = max(100, top_k * 10)
         
         # Stage 1: Prefetch with pooled vectors
-        logger.info(f"🔍 Stage 1: Prefetching {prefetch_k} candidates with pooled search")
+        logger.info(f"🔍 Stage 1: Prefetching {prefetch_k} candidates ({stage1_mode})")
         candidates = self._stage1_prefetch(
             query_np=query_np,
             top_k=prefetch_k,
             filter_obj=filter_obj,
+            stage1_mode=stage1_mode,
         )
         
         if not candidates:
@@ -202,21 +316,34 @@ class TwoStageRetriever:
         query_np: np.ndarray,
         top_k: int,
         filter_obj=None,
+        stage1_mode: str = "pooled_query_vs_tiles",
     ) -> List[Dict[str, Any]]:
-        """Stage 1: Fast prefetch with pooled vectors."""
-        # Pool query to single vector
-        query_pooled = query_np.mean(axis=0).tolist()
+        """Stage 1: Prefetch candidates."""
+        if stage1_mode == "pooled_query_vs_tiles":
+            query_vector = query_np.mean(axis=0).tolist()
+            vector_name = self.pooled_vector_name
+        elif stage1_mode == "tokens_vs_tiles":
+            query_vector = query_np.tolist()
+            vector_name = self.pooled_vector_name
+        elif stage1_mode == "pooled_query_vs_global":
+            query_vector = query_np.mean(axis=0).tolist()
+            vector_name = self.global_vector_name
+        else:
+            raise ValueError(f"Unknown stage1_mode: {stage1_mode}")
         
-        results = self.client.query_points(
-            collection_name=self.collection_name,
-            query=query_pooled,
-            using=self.pooled_vector_name,
-            query_filter=filter_obj,
-            limit=top_k,
-            with_payload=True,
-            with_vectors=False,  # Don't fetch vectors yet
-            timeout=60,
-        ).points
+        def _do_query():
+            return self.client.query_points(
+                collection_name=self.collection_name,
+                query=query_vector,
+                using=vector_name,
+                query_filter=filter_obj,
+                limit=top_k,
+                with_payload=True,
+                with_vectors=False,
+                timeout=self.request_timeout,
+            ).points
+
+        results = self._retry_call(_do_query)
         
         return [
             {
@@ -241,12 +368,16 @@ class TwoStageRetriever:
         candidate_ids = [c["id"] for c in candidates]
         
         # Retrieve points with vectors
-        points = self.client.retrieve(
-            collection_name=self.collection_name,
-            ids=candidate_ids,
-            with_payload=False,
-            with_vectors=[self.full_vector_name],
-        )
+        def _do_retrieve():
+            return self.client.retrieve(
+                collection_name=self.collection_name,
+                ids=candidate_ids,
+                with_payload=False,
+                with_vectors=[self.full_vector_name],
+                timeout=self.request_timeout,
+            )
+
+        points = self._retry_call(_do_retrieve)
         
         # Build ID to embedding map
         id_to_embedding = {}