refactor pooling options

benchmarks/__init__.py

@@ -8,4 +8,3 @@ work in Docker/Spaces environments.
 """
 
 __all__ = []
-

benchmarks/analyze_results.py

@@ -5,7 +5,7 @@ Analyze and compare benchmark results.
 Usage:
     # Compare exhaustive vs two-stage
     python analyze_results.py --results results/
-    
+
     # Compare multiple models
     python analyze_results.py --dirs results_colsmol/ results_colpali/
 """
@@ -13,7 +13,7 @@ Usage:
 import argparse
 import json
 from pathlib import Path
-from typing import Dict, List, Any
+from typing import Dict
 
 import numpy as np
 
@@ -24,12 +24,12 @@ def load_all_results(results_dir: Path) -> Dict[str, Dict]:
     for f in results_dir.glob("*.json"):
         with open(f) as fp:
             data = json.load(fp)
-        
+
         # Key by dataset + method
         dataset = data.get("dataset", f.stem).split("/")[-1]
         method = "two_stage" if data.get("two_stage") else "exhaustive"
         key = f"{dataset}_{method}"
-        
+
         results[key] = {
             "dataset": dataset,
             "method": method,
@@ -41,7 +41,7 @@ def load_all_results(results_dir: Path) -> Dict[str, Dict]:
 
 def compare_methods(results: Dict[str, Dict]) -> None:
     """Compare exhaustive vs two-stage on same datasets."""
-    
+
     # Group by dataset
     datasets = {}
     for key, data in results.items():
@@ -49,45 +49,47 @@ def compare_methods(results: Dict[str, Dict]) -> None:
         if ds not in datasets:
             datasets[ds] = {}
         datasets[ds][data["method"]] = data
-    
+
     print("\n" + "=" * 80)
     print("EXHAUSTIVE vs TWO-STAGE COMPARISON")
     print("=" * 80)
-    
+
     print(f"\n{'Dataset':<30} {'Method':<12} {'NDCG@10':>10} {'MRR@10':>10} {'Time(ms)':>10}")
     print("-" * 72)
-    
+
     improvements = []
     speedups = []
-    
+
     for dataset, methods in sorted(datasets.items()):
         for method in ["exhaustive", "two_stage"]:
             if method in methods:
                 m = methods[method]
                 time_ms = m.get("avg_search_time_ms", 0)
-                print(f"{dataset:<30} {method:<12} {m.get('ndcg@10', 0):>10.4f} {m.get('mrr@10', 0):>10.4f} {time_ms:>10.2f}")
-        
+                print(
+                    f"{dataset:<30} {method:<12} {m.get('ndcg@10', 0):>10.4f} {m.get('mrr@10', 0):>10.4f} {time_ms:>10.2f}"
+                )
+
         # Calculate improvement
         if "exhaustive" in methods and "two_stage" in methods:
             ex = methods["exhaustive"]
             ts = methods["two_stage"]
-            
+
             ndcg_diff = ts.get("ndcg@10", 0) - ex.get("ndcg@10", 0)
             improvements.append(ndcg_diff)
-            
+
             ex_time = ex.get("avg_search_time_ms", 1)
             ts_time = ts.get("avg_search_time_ms", 1)
             if ts_time > 0:
                 speedups.append(ex_time / ts_time)
-        
+
         print()
-    
+
     # Summary
     if improvements:
         print("-" * 72)
         print(f"Average NDCG@10 difference (two_stage - exhaustive): {np.mean(improvements):+.4f}")
         print(f"Retention rate: {100 * (1 + np.mean(improvements)):.1f}%")
-    
+
     if speedups:
         print(f"Average speedup: {np.mean(speedups):.1f}x")
 
@@ -106,7 +108,7 @@ def print_leaderboard(results: Dict[str, Dict]) -> None:
     print("\n" + "=" * 80)
     print("LEADERBOARD FORMAT")
     print("=" * 80)
-    
+
     # Best result per dataset
     best = {}
     for key, data in results.items():
@@ -114,44 +116,32 @@ def print_leaderboard(results: Dict[str, Dict]) -> None:
         ndcg = data.get("ndcg@10", 0)
         if ds not in best or ndcg > best[ds].get("ndcg@10", 0):
             best[ds] = data
-    
+
     # Compute average
     ndcg_scores = [d.get("ndcg@10", 0) for d in best.values()]
     avg = sum(ndcg_scores) / len(ndcg_scores) if ndcg_scores else 0
-    
+
     print(f"\nModel: {list(results.values())[0].get('model', 'unknown')}")
     print(f"\n{'Dataset':<35} {'NDCG@10':>10}")
     print("-" * 45)
-    
+
     for ds, data in sorted(best.items()):
         method_tag = " (2-stage)" if data.get("method") == "two_stage" else ""
         print(f"{ds + method_tag:<35} {data.get('ndcg@10', 0):>10.4f}")
-    
+
     print("-" * 45)
     print(f"{'AVERAGE':<35} {avg:>10.4f}")
 
 
 def main():
     parser = argparse.ArgumentParser(description="Analyze benchmark results")
-    parser.add_argument(
-        "--results", type=str, default="results",
-        help="Results directory"
-    )
-    parser.add_argument(
-        "--dirs", nargs="+",
-        help="Multiple result directories to compare"
-    )
-    parser.add_argument(
-        "--compare", action="store_true",
-        help="Compare exhaustive vs two-stage"
-    )
-    parser.add_argument(
-        "--leaderboard", action="store_true",
-        help="Print in leaderboard format"
-    )
-    
+    parser.add_argument("--results", type=str, default="results", help="Results directory")
+    parser.add_argument("--dirs", nargs="+", help="Multiple result directories to compare")
+    parser.add_argument("--compare", action="store_true", help="Compare exhaustive vs two-stage")
+    parser.add_argument("--leaderboard", action="store_true", help="Print in leaderboard format")
+
     args = parser.parse_args()
-    
+
     if args.dirs:
         # Compare multiple directories
         all_results = {}
@@ -162,26 +152,19 @@ def main():
         results = all_results
     else:
         results = load_all_results(Path(args.results))
-    
+
     if not results:
-        print(f"❌ No results found")
+        print("❌ No results found")
         return
-    
+
     print(f"📊 Loaded {len(results)} result files")
-    
+
     if args.compare or not args.leaderboard:
         compare_methods(results)
-    
+
     if args.leaderboard or not args.compare:
         print_leaderboard(results)
 
 
 if __name__ == "__main__":
     main()
-
-
-
-
-
-
-

benchmarks/prepare_submission.py

@@ -11,14 +11,14 @@ Usage:
 
 import argparse
 import json
-from pathlib import Path
 from datetime import datetime
-from typing import Dict, Any, Optional
+from pathlib import Path
+from typing import Any, Dict, Optional
 
 # ViDoRe leaderboard expected datasets
 VIDORE_DATASETS = {
     "docvqa_test_subsampled": "DocVQA",
-    "infovqa_test_subsampled": "InfoVQA", 
+    "infovqa_test_subsampled": "InfoVQA",
     "tabfquad_test_subsampled": "TabFQuAD",
     "tatdqa_test": "TAT-DQA",
     "arxivqa_test_subsampled": "ArXivQA",
@@ -29,14 +29,14 @@ VIDORE_DATASETS = {
 def load_results(results_dir: Path) -> Dict[str, Dict[str, float]]:
     """Load all result JSON files from directory."""
     results = {}
-    
+
     for json_file in results_dir.glob("*.json"):
         with open(json_file) as f:
             data = json.load(f)
-        
+
         dataset = data.get("dataset", json_file.stem)
         dataset_short = dataset.split("/")[-1].replace("_twostage", "")
-        
+
         results[dataset_short] = {
             "ndcg@5": data["metrics"].get("ndcg@5", 0),
             "ndcg@10": data["metrics"].get("ndcg@10", 0),
@@ -46,7 +46,7 @@ def load_results(results_dir: Path) -> Dict[str, Dict[str, float]]:
             "two_stage": data.get("two_stage", False),
             "model": data.get("model", "unknown"),
         }
-    
+
     return results
 
 
@@ -57,11 +57,11 @@ def format_submission(
     description: Optional[str] = None,
 ) -> Dict[str, Any]:
     """Format results for ViDoRe leaderboard submission."""
-    
+
     # Calculate average scores
     ndcg10_scores = [r["ndcg@10"] for r in results.values()]
     avg_ndcg10 = sum(ndcg10_scores) / len(ndcg10_scores) if ndcg10_scores else 0
-    
+
     submission = {
         "model_name": model_name,
         "model_url": model_url or "",
@@ -70,7 +70,7 @@ def format_submission(
         "average_ndcg@10": avg_ndcg10,
         "results": {},
     }
-    
+
     # Add per-dataset results
     for dataset_short, metrics in results.items():
         display_name = VIDORE_DATASETS.get(dataset_short, dataset_short)
@@ -79,7 +79,7 @@ def format_submission(
             "ndcg@10": metrics["ndcg@10"],
             "mrr@10": metrics["mrr@10"],
         }
-    
+
     return submission
 
 
@@ -88,14 +88,16 @@ def print_summary(results: Dict[str, Dict], submission: Dict[str, Any]):
     print("\n" + "=" * 70)
     print(f"MODEL: {submission['model_name']}")
     print("=" * 70)
-    
+
     print(f"\n{'Dataset':<25} {'NDCG@5':>10} {'NDCG@10':>10} {'MRR@10':>10}")
     print("-" * 55)
-    
+
     for dataset, metrics in results.items():
         display = VIDORE_DATASETS.get(dataset, dataset)[:24]
-        print(f"{display:<25} {metrics['ndcg@5']:>10.4f} {metrics['ndcg@10']:>10.4f} {metrics['mrr@10']:>10.4f}")
-    
+        print(
+            f"{display:<25} {metrics['ndcg@5']:>10.4f} {metrics['ndcg@10']:>10.4f} {metrics['mrr@10']:>10.4f}"
+        )
+
     print("-" * 55)
     print(f"{'AVERAGE':<25} {'':<10} {submission['average_ndcg@10']:>10.4f}")
     print("=" * 70)
@@ -108,14 +110,14 @@ def upload_to_huggingface(submission: Dict[str, Any], repo_id: str = "vidore/res
     except ImportError:
         print("Install huggingface_hub: pip install huggingface_hub")
         return False
-    
+
     api = HfApi()
-    
+
     # Save to temp file
     temp_file = Path(f"/tmp/submission_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json")
     with open(temp_file, "w") as f:
         json.dump(submission, f, indent=2)
-    
+
     try:
         api.upload_file(
             path_or_fileobj=str(temp_file),
@@ -133,45 +135,33 @@ def upload_to_huggingface(submission: Dict[str, Any], repo_id: str = "vidore/res
 def main():
     parser = argparse.ArgumentParser(description="Prepare ViDoRe submission")
     parser.add_argument(
-        "--results", type=str, default="results",
-        help="Directory with result JSON files"
-    )
-    parser.add_argument(
-        "--output", type=str, default="submission.json",
-        help="Output submission file"
+        "--results", type=str, default="results", help="Directory with result JSON files"
     )
     parser.add_argument(
-        "--model-name", type=str, default="visual-rag-toolkit",
-        help="Model name for leaderboard"
+        "--output", type=str, default="submission.json", help="Output submission file"
     )
     parser.add_argument(
-        "--model-url", type=str,
-        help="URL to model/paper"
+        "--model-name", type=str, default="visual-rag-toolkit", help="Model name for leaderboard"
     )
-    parser.add_argument(
-        "--description", type=str,
-        help="Model description"
-    )
-    parser.add_argument(
-        "--upload", action="store_true",
-        help="Upload to HuggingFace"
-    )
-    
+    parser.add_argument("--model-url", type=str, help="URL to model/paper")
+    parser.add_argument("--description", type=str, help="Model description")
+    parser.add_argument("--upload", action="store_true", help="Upload to HuggingFace")
+
     args = parser.parse_args()
-    
+
     results_dir = Path(args.results)
     if not results_dir.exists():
         print(f"❌ Results directory not found: {results_dir}")
         return
-    
+
     # Load results
     results = load_results(results_dir)
     if not results:
         print(f"❌ No result files found in {results_dir}")
         return
-    
+
     print(f"📊 Found {len(results)} dataset results")
-    
+
     # Format submission
     submission = format_submission(
         results,
@@ -179,16 +169,16 @@ def main():
         model_url=args.model_url,
         description=args.description,
     )
-    
+
     # Print summary
     print_summary(results, submission)
-    
+
     # Save
     output_path = Path(args.output)
     with open(output_path, "w") as f:
         json.dump(submission, f, indent=2)
     print(f"\n💾 Saved to: {output_path}")
-    
+
     # Upload if requested
     if args.upload:
         upload_to_huggingface(submission)
@@ -196,10 +186,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-
-
-
-
-

benchmarks/quick_test.py

@@ -14,12 +14,12 @@ Usage:
     python quick_test.py --samples 500 --skip-exhaustive  # Faster
 """
 
-import sys
-import time
 import argparse
 import logging
+import sys
+import time
 from pathlib import Path
-from typing import List, Dict, Any
+from typing import Any, Dict, List
 
 # Add parent directory to Python path (so we can import visual_rag)
 # This allows running the script directly without pip install
@@ -28,57 +28,60 @@ _parent_dir = _script_dir.parent
 if str(_parent_dir) not in sys.path:
     sys.path.insert(0, str(_parent_dir))
 
-import numpy as np
-from tqdm import tqdm
+import numpy as np  # noqa: E402
+from tqdm import tqdm  # noqa: E402
 
 # Visual RAG imports (now works without pip install)
-from visual_rag.embedding import VisualEmbedder
-from visual_rag.embedding.pooling import (
-    tile_level_mean_pooling,
+from visual_rag.embedding import VisualEmbedder  # noqa: E402
+from visual_rag.embedding.pooling import (  # noqa: E402
     compute_maxsim_score,
+    tile_level_mean_pooling,
 )
 
 # Optional: datasets for ViDoRe
 try:
     from datasets import load_dataset as hf_load_dataset
+
     HAS_DATASETS = True
 except ImportError:
     HAS_DATASETS = False
 
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
 
 def load_vidore_sample(num_samples: int = 100) -> List[Dict]:
     """
     Load sample from ViDoRe DocVQA with ground truth.
-    
+
     Each sample has a query and its relevant document (1:1 mapping).
     This allows computing retrieval metrics.
     """
     if not HAS_DATASETS:
         logger.error("Install datasets: pip install datasets")
         sys.exit(1)
-    
+
     logger.info(f"📥 Loading {num_samples} samples from ViDoRe DocVQA...")
-    
+
     ds = hf_load_dataset("vidore/docvqa_test_subsampled", split="test")
-    
+
     samples = []
     for i, example in enumerate(ds):
         if i >= num_samples:
             break
-        
-        samples.append({
-            "id": i,
-            "doc_id": f"doc_{i}",
-            "query_id": f"q_{i}",
-            "image": example.get("image", example.get("page_image")),
-            "query": example.get("query", example.get("question", "")),
-            # Ground truth: query i is relevant to doc i
-            "relevant_doc": f"doc_{i}",
-        })
-    
+
+        samples.append(
+            {
+                "id": i,
+                "doc_id": f"doc_{i}",
+                "query_id": f"q_{i}",
+                "image": example.get("image", example.get("page_image")),
+                "query": example.get("query", example.get("question", "")),
+                # Ground truth: query i is relevant to doc i
+                "relevant_doc": f"doc_{i}",
+            }
+        )
+
     logger.info(f"✅ Loaded {len(samples)} samples with ground truth")
     return samples
 
@@ -90,60 +93,58 @@ def embed_all(
     """Embed all documents and queries."""
     logger.info(f"\n🤖 Loading model: {model_name}")
     embedder = VisualEmbedder(model_name=model_name)
-    
+
     images = [s["image"] for s in samples]
     queries = [s["query"] for s in samples if s["query"]]
-    
+
     # Embed images
     logger.info(f"🎨 Embedding {len(images)} documents...")
     start_time = time.time()
-    
-    embeddings, token_infos = embedder.embed_images(
-        images, batch_size=4, return_token_info=True
-    )
-    
+
+    embeddings, token_infos = embedder.embed_images(images, batch_size=4, return_token_info=True)
+
     doc_embed_time = time.time() - start_time
     logger.info(f"   Time: {doc_embed_time:.2f}s ({doc_embed_time/len(images)*1000:.1f}ms/doc)")
-    
+
     # Process embeddings: extract visual tokens + tile-level pooling
     doc_data = {}
     for i, (emb, token_info) in enumerate(zip(embeddings, token_infos)):
-        if hasattr(emb, 'cpu'):
+        if hasattr(emb, "cpu"):
             emb = emb.cpu()
-        emb_np = emb.numpy() if hasattr(emb, 'numpy') else np.array(emb)
-        
+        emb_np = emb.numpy() if hasattr(emb, "numpy") else np.array(emb)
+
         # Extract visual tokens only (filter special tokens)
         visual_indices = token_info["visual_token_indices"]
         visual_emb = emb_np[visual_indices].astype(np.float32)
-        
+
         # Tile-level pooling
         n_rows = token_info.get("n_rows", 4)
         n_cols = token_info.get("n_cols", 3)
         num_tiles = n_rows * n_cols + 1 if n_rows and n_cols else 13
-        
+
         tile_pooled = tile_level_mean_pooling(visual_emb, num_tiles, patches_per_tile=64)
-        
+
         doc_data[f"doc_{i}"] = {
             "embedding": visual_emb,
             "pooled": tile_pooled,
             "num_visual_tokens": len(visual_indices),
             "num_tiles": tile_pooled.shape[0],
         }
-    
+
     # Embed queries
     logger.info(f"🔍 Embedding {len(queries)} queries...")
     start_time = time.time()
-    
+
     query_data = {}
     for i, query in enumerate(tqdm(queries, desc="Queries")):
         q_emb = embedder.embed_query(query)
-        if hasattr(q_emb, 'cpu'):
+        if hasattr(q_emb, "cpu"):
             q_emb = q_emb.cpu()
-        q_np = q_emb.numpy() if hasattr(q_emb, 'numpy') else np.array(q_emb)
+        q_np = q_emb.numpy() if hasattr(q_emb, "numpy") else np.array(q_emb)
         query_data[f"q_{i}"] = q_np.astype(np.float32)
-    
+
     query_embed_time = time.time() - start_time
-    
+
     return {
         "docs": doc_data,
         "queries": query_data,
@@ -160,7 +161,7 @@ def search_exhaustive(query_emb: np.ndarray, docs: Dict, top_k: int = 10) -> Lis
     for doc_id, doc in docs.items():
         score = compute_maxsim_score(query_emb, doc["embedding"])
         scores.append({"id": doc_id, "score": score})
-    
+
     scores.sort(key=lambda x: x["score"], reverse=True)
     return scores[:top_k]
 
@@ -173,14 +174,14 @@ def search_two_stage(
 ) -> List[Dict]:
     """
     Two-stage retrieval with tile-level pooling.
-    
+
     Stage 1: Fast prefetch using tile-pooled vectors
     Stage 2: Exact MaxSim reranking on candidates
     """
     # Stage 1: Tile-level pooled search
     query_pooled = query_emb.mean(axis=0)
     query_pooled = query_pooled / (np.linalg.norm(query_pooled) + 1e-8)
-    
+
     stage1_scores = []
     for doc_id, doc in docs.items():
         doc_pooled = doc["pooled"]
@@ -188,17 +189,19 @@ def search_two_stage(
         tile_sims = np.dot(doc_norm, query_pooled)
         score = float(tile_sims.max())
         stage1_scores.append({"id": doc_id, "score": score})
-    
+
     stage1_scores.sort(key=lambda x: x["score"], reverse=True)
     candidates = stage1_scores[:prefetch_k]
-    
+
     # Stage 2: Exact MaxSim on candidates
     reranked = []
     for cand in candidates:
         doc_id = cand["id"]
         score = compute_maxsim_score(query_emb, docs[doc_id]["embedding"])
-        reranked.append({"id": doc_id, "score": score, "stage1_rank": stage1_scores.index(cand) + 1})
-    
+        reranked.append(
+            {"id": doc_id, "score": score, "stage1_rank": stage1_scores.index(cand) + 1}
+        )
+
     reranked.sort(key=lambda x: x["score"], reverse=True)
     return reranked[:top_k]
 
@@ -210,54 +213,54 @@ def compute_metrics(
 ) -> Dict[str, float]:
     """
     Compute retrieval metrics.
-    
+
     Since ViDoRe has 1:1 query-doc mapping (1 relevant doc per query):
     - Recall@K (Hit Rate): Is the relevant doc in top-K? (0 or 1)
-    - Precision@K: (# relevant in top-K) / K 
+    - Precision@K: (# relevant in top-K) / K
     - MRR@K: 1/rank if found in top-K, else 0
     - NDCG@K: DCG / IDCG with binary relevance
     """
     metrics = {}
-    
+
     # Also track per-query ranks for analysis
     all_ranks = []
-    
+
     for k in k_values:
         recalls = []
         precisions = []
         mrrs = []
         ndcgs = []
-        
+
         for sample in samples:
             query_id = sample["query_id"]
             relevant_doc = sample["relevant_doc"]
-            
+
             if query_id not in results:
                 continue
-            
+
             ranking = results[query_id][:k]
             ranked_ids = [r["id"] for r in ranking]
-            
+
             # Find rank of relevant doc (1-indexed, 0 if not found)
             rank = 0
             for i, doc_id in enumerate(ranked_ids):
                 if doc_id == relevant_doc:
                     rank = i + 1
                     break
-            
+
             # Recall@K (Hit Rate): 1 if found in top-K
             found = 1.0 if rank > 0 else 0.0
             recalls.append(found)
-            
+
             # Precision@K: (# relevant found) / K
             # With 1 relevant doc: 1/K if found, 0 otherwise
             precision = found / k
             precisions.append(precision)
-            
+
             # MRR@K: 1/rank if found
             mrr = 1.0 / rank if rank > 0 else 0.0
             mrrs.append(mrr)
-            
+
             # NDCG@K (binary relevance)
             # DCG = 1/log2(rank+1) if found, 0 otherwise
             # IDCG = 1/log2(2) = 1 (best case: relevant at rank 1)
@@ -265,7 +268,7 @@ def compute_metrics(
             idcg = 1.0
             ndcg = dcg / idcg
             ndcgs.append(ndcg)
-            
+
             # Track actual rank for analysis (only for k=10)
             if k == max(k_values):
                 full_ranking = results[query_id]
@@ -275,19 +278,19 @@ def compute_metrics(
                         full_rank = i + 1
                         break
                 all_ranks.append(full_rank)
-        
+
         metrics[f"Recall@{k}"] = np.mean(recalls)
         metrics[f"P@{k}"] = np.mean(precisions)
         metrics[f"MRR@{k}"] = np.mean(mrrs)
         metrics[f"NDCG@{k}"] = np.mean(ndcgs)
-    
+
     # Add summary stats
     if all_ranks:
         found_ranks = [r for r in all_ranks if r > 0]
-        metrics["avg_rank"] = np.mean(found_ranks) if found_ranks else float('inf')
-        metrics["median_rank"] = np.median(found_ranks) if found_ranks else float('inf')
+        metrics["avg_rank"] = np.mean(found_ranks) if found_ranks else float("inf")
+        metrics["median_rank"] = np.median(found_ranks) if found_ranks else float("inf")
         metrics["not_found"] = sum(1 for r in all_ranks if r == 0)
-    
+
     return metrics
 
 
@@ -302,67 +305,71 @@ def run_benchmark(
     queries = data["queries"]
     samples = data["samples"]
     num_docs = len(docs)
-    
+
     # Auto-set prefetch_k to be meaningful (default: 20, or 20% of docs if >100 docs)
     if prefetch_k is None:
         if num_docs <= 100:
             prefetch_k = 20  # Default: prefetch 20, rerank to top-10
         else:
             prefetch_k = max(20, min(100, int(num_docs * 0.2)))  # 20% for larger collections
-    
+
     # Ensure prefetch_k < num_docs for meaningful two-stage comparison
     if prefetch_k >= num_docs:
         logger.warning(f"⚠️  prefetch_k={prefetch_k} >= num_docs={num_docs}")
-        logger.warning(f"   Two-stage will fetch ALL docs (same as exhaustive)")
+        logger.warning("   Two-stage will fetch ALL docs (same as exhaustive)")
         logger.warning(f"   Use --samples > {prefetch_k * 3} for meaningful comparison")
-    
+
     logger.info(f"📊 Benchmark config: {num_docs} docs, prefetch_k={prefetch_k}, top_k={top_k}")
     logger.info(f"   (Both methods return top-{top_k} results - realistic retrieval scenario)")
-    
+
     results = {}
-    
+
     # Two-stage retrieval (NOVEL)
-    logger.info(f"\n🔬 Running Two-Stage retrieval (prefetch top-{prefetch_k}, rerank to top-{top_k})...")
+    logger.info(
+        f"\n🔬 Running Two-Stage retrieval (prefetch top-{prefetch_k}, rerank to top-{top_k})..."
+    )
     two_stage_results = {}
     two_stage_times = []
-    
+
     for sample in tqdm(samples, desc="Two-Stage"):
         query_id = sample["query_id"]
         query_emb = queries[query_id]
-        
+
         start = time.time()
         ranking = search_two_stage(query_emb, docs, prefetch_k=prefetch_k, top_k=top_k)
         two_stage_times.append(time.time() - start)
-        
+
         two_stage_results[query_id] = ranking
-    
+
     two_stage_metrics = compute_metrics(two_stage_results, samples)
     two_stage_metrics["avg_time_ms"] = np.mean(two_stage_times) * 1000
     two_stage_metrics["prefetch_k"] = prefetch_k
     two_stage_metrics["top_k"] = top_k
     results["two_stage"] = two_stage_metrics
-    
+
     # Exhaustive search (baseline)
     if not skip_exhaustive:
-        logger.info(f"🔬 Running Exhaustive MaxSim (searches ALL {num_docs} docs, returns top-{top_k})...")
+        logger.info(
+            f"🔬 Running Exhaustive MaxSim (searches ALL {num_docs} docs, returns top-{top_k})..."
+        )
         exhaustive_results = {}
         exhaustive_times = []
-        
+
         for sample in tqdm(samples, desc="Exhaustive"):
             query_id = sample["query_id"]
             query_emb = queries[query_id]
-            
+
             start = time.time()
             ranking = search_exhaustive(query_emb, docs, top_k=top_k)
             exhaustive_times.append(time.time() - start)
-            
+
             exhaustive_results[query_id] = ranking
-        
+
         exhaustive_metrics = compute_metrics(exhaustive_results, samples)
         exhaustive_metrics["avg_time_ms"] = np.mean(exhaustive_times) * 1000
         exhaustive_metrics["top_k"] = top_k
         results["exhaustive"] = exhaustive_metrics
-    
+
     return results
 
 
@@ -371,88 +378,98 @@ def print_results(data: Dict, benchmark_results: Dict, show_precision: bool = Fa
     print("\n" + "=" * 80)
     print("📊 BENCHMARK RESULTS")
     print("=" * 80)
-    
-    num_docs = len(data['docs'])
+
+    num_docs = len(data["docs"])
     print(f"\n🤖 Model: {data['model']}")
     print(f"📄 Documents: {num_docs}")
     print(f"🔍 Queries: {len(data['queries'])}")
-    
+
     # Embedding stats
-    sample_doc = list(data['docs'].values())[0]
-    print(f"\n📏 Embedding (after visual token filtering):")
+    sample_doc = list(data["docs"].values())[0]
+    print("\n📏 Embedding (after visual token filtering):")
     print(f"   Visual tokens per doc: {sample_doc['num_visual_tokens']}")
     print(f"   Tile-pooled vectors: {sample_doc['num_tiles']}")
-    
+
     if "two_stage" in benchmark_results:
         prefetch_k = benchmark_results["two_stage"].get("prefetch_k", "?")
         print(f"   Two-stage prefetch_k: {prefetch_k} (of {num_docs} docs)")
-    
+
     # Method labels - clearer naming
     def get_label(method):
         if method == "two_stage":
             return "Pooled+Rerank"  # Tile-pooled prefetch + MaxSim rerank
         else:
-            return "Full MaxSim"    # Exhaustive MaxSim on all docs
-    
+            return "Full MaxSim"  # Exhaustive MaxSim on all docs
+
     # Recall / Hit Rate table
-    print(f"\n🎯 RECALL (Hit Rate) @ K:")
+    print("\n🎯 RECALL (Hit Rate) @ K:")
     print(f"   {'Method':<20} {'@1':>8} {'@3':>8} {'@5':>8} {'@7':>8} {'@10':>8}")
     print(f"   {'-'*60}")
-    
+
     for method, metrics in benchmark_results.items():
-        print(f"   {get_label(method):<20} "
-              f"{metrics.get('Recall@1', 0):>8.3f} "
-              f"{metrics.get('Recall@3', 0):>8.3f} "
-              f"{metrics.get('Recall@5', 0):>8.3f} "
-              f"{metrics.get('Recall@7', 0):>8.3f} "
-              f"{metrics.get('Recall@10', 0):>8.3f}")
-    
+        print(
+            f"   {get_label(method):<20} "
+            f"{metrics.get('Recall@1', 0):>8.3f} "
+            f"{metrics.get('Recall@3', 0):>8.3f} "
+            f"{metrics.get('Recall@5', 0):>8.3f} "
+            f"{metrics.get('Recall@7', 0):>8.3f} "
+            f"{metrics.get('Recall@10', 0):>8.3f}"
+        )
+
     # Precision table (optional)
     if show_precision:
-        print(f"\n📐 PRECISION @ K:")
+        print("\n📐 PRECISION @ K:")
         print(f"   {'Method':<20} {'@1':>8} {'@3':>8} {'@5':>8} {'@7':>8} {'@10':>8}")
         print(f"   {'-'*60}")
-        
+
         for method, metrics in benchmark_results.items():
-            print(f"   {get_label(method):<20} "
-                  f"{metrics.get('P@1', 0):>8.3f} "
-                  f"{metrics.get('P@3', 0):>8.3f} "
-                  f"{metrics.get('P@5', 0):>8.3f} "
-                  f"{metrics.get('P@7', 0):>8.3f} "
-                  f"{metrics.get('P@10', 0):>8.3f}")
-    
+            print(
+                f"   {get_label(method):<20} "
+                f"{metrics.get('P@1', 0):>8.3f} "
+                f"{metrics.get('P@3', 0):>8.3f} "
+                f"{metrics.get('P@5', 0):>8.3f} "
+                f"{metrics.get('P@7', 0):>8.3f} "
+                f"{metrics.get('P@10', 0):>8.3f}"
+            )
+
     # NDCG table
-    print(f"\n📈 NDCG @ K:")
+    print("\n📈 NDCG @ K:")
     print(f"   {'Method':<20} {'@1':>8} {'@3':>8} {'@5':>8} {'@7':>8} {'@10':>8}")
     print(f"   {'-'*60}")
-    
+
     for method, metrics in benchmark_results.items():
-        print(f"   {get_label(method):<20} "
-              f"{metrics.get('NDCG@1', 0):>8.3f} "
-              f"{metrics.get('NDCG@3', 0):>8.3f} "
-              f"{metrics.get('NDCG@5', 0):>8.3f} "
-              f"{metrics.get('NDCG@7', 0):>8.3f} "
-              f"{metrics.get('NDCG@10', 0):>8.3f}")
-    
+        print(
+            f"   {get_label(method):<20} "
+            f"{metrics.get('NDCG@1', 0):>8.3f} "
+            f"{metrics.get('NDCG@3', 0):>8.3f} "
+            f"{metrics.get('NDCG@5', 0):>8.3f} "
+            f"{metrics.get('NDCG@7', 0):>8.3f} "
+            f"{metrics.get('NDCG@10', 0):>8.3f}"
+        )
+
     # MRR table
-    print(f"\n🔍 MRR @ K:")
+    print("\n🔍 MRR @ K:")
     print(f"   {'Method':<20} {'@1':>8} {'@3':>8} {'@5':>8} {'@7':>8} {'@10':>8}")
     print(f"   {'-'*60}")
-    
+
     for method, metrics in benchmark_results.items():
-        print(f"   {get_label(method):<20} "
-              f"{metrics.get('MRR@1', 0):>8.3f} "
-              f"{metrics.get('MRR@3', 0):>8.3f} "
-              f"{metrics.get('MRR@5', 0):>8.3f} "
-              f"{metrics.get('MRR@7', 0):>8.3f} "
-              f"{metrics.get('MRR@10', 0):>8.3f}")
-    
+        print(
+            f"   {get_label(method):<20} "
+            f"{metrics.get('MRR@1', 0):>8.3f} "
+            f"{metrics.get('MRR@3', 0):>8.3f} "
+            f"{metrics.get('MRR@5', 0):>8.3f} "
+            f"{metrics.get('MRR@7', 0):>8.3f} "
+            f"{metrics.get('MRR@10', 0):>8.3f}"
+        )
+
     # Speed comparison
-    top_k = benchmark_results.get("two_stage", benchmark_results.get("exhaustive", {})).get("top_k", 10)
+    top_k = benchmark_results.get("two_stage", benchmark_results.get("exhaustive", {})).get(
+        "top_k", 10
+    )
     print(f"\n⏱️  SPEED (both return top-{top_k} results):")
     print(f"   {'Method':<20} {'Time (ms)':>12} {'Docs searched':>15}")
     print(f"   {'-'*50}")
-    
+
     for method, metrics in benchmark_results.items():
         if method == "two_stage":
             searched = metrics.get("prefetch_k", "?")
@@ -461,45 +478,53 @@ def print_results(data: Dict, benchmark_results: Dict, show_precision: bool = Fa
             searched = num_docs
             label = f"{searched} (all)"
         print(f"   {get_label(method):<20} {metrics.get('avg_time_ms', 0):>12.2f} {label:>15}")
-    
+
     # Comparison summary
     if "exhaustive" in benchmark_results and "two_stage" in benchmark_results:
         ex = benchmark_results["exhaustive"]
         ts = benchmark_results["two_stage"]
-        
-        print(f"\n💡 POOLED+RERANK vs FULL MAXSIM:")
-        
+
+        print("\n💡 POOLED+RERANK vs FULL MAXSIM:")
+
         for k in [1, 5, 10]:
             ex_recall = ex.get(f"Recall@{k}", 0)
             ts_recall = ts.get(f"Recall@{k}", 0)
             if ex_recall > 0:
                 retention = ts_recall / ex_recall * 100
-                print(f"   • Recall@{k} retention: {retention:.1f}% ({ts_recall:.3f} vs {ex_recall:.3f})")
-        
+                print(
+                    f"   • Recall@{k} retention: {retention:.1f}% ({ts_recall:.3f} vs {ex_recall:.3f})"
+                )
+
         speedup = ex["avg_time_ms"] / ts["avg_time_ms"] if ts["avg_time_ms"] > 0 else 0
         print(f"   • Speedup: {speedup:.1f}x")
-        
+
         # Rank stats with explanation
         if "avg_rank" in ts:
             prefetch_k = ts.get("prefetch_k", "?")
             top_k = ts.get("top_k", 10)
             not_found = ts.get("not_found", 0)
             total = len(data["queries"])
-            
-            print(f"\n📊 POOLED+RERANK STATISTICS:")
-            print(f"   Stage-1 (pooled prefetch):")
+
+            print("\n📊 POOLED+RERANK STATISTICS:")
+            print("   Stage-1 (pooled prefetch):")
             print(f"      • Searches top-{prefetch_k} candidates using tile-pooled vectors")
-            print(f"      • {total - not_found}/{total} queries ({100 - not_found/total*100:.1f}%) had relevant doc in prefetch")
-            print(f"      • {not_found}/{total} queries ({not_found/total*100:.1f}%) missed (relevant doc ranked >{prefetch_k})")
-            print(f"   Stage-2 (MaxSim reranking):")
-            print(f"      • Reranks prefetch candidates with exact MaxSim")
+            print(
+                f"      • {total - not_found}/{total} queries ({100 - not_found/total*100:.1f}%) had relevant doc in prefetch"
+            )
+            print(
+                f"      • {not_found}/{total} queries ({not_found/total*100:.1f}%) missed (relevant doc ranked >{prefetch_k})"
+            )
+            print("   Stage-2 (MaxSim reranking):")
+            print("      • Reranks prefetch candidates with exact MaxSim")
             print(f"      • Returns final top-{top_k} results")
-            if ts['avg_rank'] < float('inf'):
+            if ts["avg_rank"] < float("inf"):
                 print(f"      • Avg rank of relevant doc (when found): {ts['avg_rank']:.1f}")
                 print(f"      • Median rank: {ts['median_rank']:.1f}")
             print(f"\n   💡 The {not_found/total*100:.1f}% miss rate is for stage-1 prefetch.")
-            print(f"      Final Recall@{top_k} shows how many relevant docs ARE in top-{top_k} results.")
-    
+            print(
+                f"      Final Recall@{top_k} shows how many relevant docs ARE in top-{top_k} results."
+            )
+
     print("\n" + "=" * 80)
     print("✅ Benchmark complete!")
 
@@ -509,47 +534,48 @@ def main():
         description="Quick benchmark for visual-rag-toolkit",
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
+    parser.add_argument("--samples", type=int, default=100, help="Number of samples (default: 100)")
     parser.add_argument(
-        "--samples", type=int, default=100,
-        help="Number of samples (default: 100)"
-    )
-    parser.add_argument(
-        "--model", type=str, default="vidore/colSmol-500M",
-        help="Model: vidore/colSmol-500M (default), vidore/colpali-v1.3"
+        "--model",
+        type=str,
+        default="vidore/colSmol-500M",
+        help="Model: vidore/colSmol-500M (default), vidore/colpali-v1.3",
     )
     parser.add_argument(
-        "--prefetch-k", type=int, default=None,
-        help="Stage 1 candidates for two-stage (default: 20 for <=100 docs, auto for larger)"
+        "--prefetch-k",
+        type=int,
+        default=None,
+        help="Stage 1 candidates for two-stage (default: 20 for <=100 docs, auto for larger)",
     )
     parser.add_argument(
-        "--skip-exhaustive", action="store_true",
-        help="Skip exhaustive baseline (faster)"
+        "--skip-exhaustive", action="store_true", help="Skip exhaustive baseline (faster)"
     )
     parser.add_argument(
-        "--show-precision", action="store_true",
-        help="Show Precision@K metrics (hidden by default)"
+        "--show-precision", action="store_true", help="Show Precision@K metrics (hidden by default)"
     )
     parser.add_argument(
-        "--top-k", type=int, default=10,
-        help="Number of results to return (default: 10, realistic retrieval scenario)"
+        "--top-k",
+        type=int,
+        default=10,
+        help="Number of results to return (default: 10, realistic retrieval scenario)",
     )
-    
+
     args = parser.parse_args()
-    
+
     print("\n" + "=" * 70)
     print("🧪 VISUAL RAG TOOLKIT - RETRIEVAL BENCHMARK")
     print("=" * 70)
-    
+
     # Load samples
     samples = load_vidore_sample(args.samples)
-    
+
     if not samples:
         logger.error("No samples loaded!")
         sys.exit(1)
-    
+
     # Embed all
     data = embed_all(samples, args.model)
-    
+
     # Run benchmark
     benchmark_results = run_benchmark(
         data,
@@ -557,7 +583,7 @@ def main():
         prefetch_k=args.prefetch_k,
         top_k=args.top_k,
     )
-    
+
     # Print results
     print_results(data, benchmark_results, show_precision=args.show_precision)
 

benchmarks/run_vidore.py

@@ -17,10 +17,10 @@ Usage:
 
 import argparse
 import json
-import time
 import logging
+import time
 from pathlib import Path
-from typing import List, Dict, Any, Optional
+from typing import Any, Dict, List, Optional
 
 import numpy as np
 from tqdm import tqdm
@@ -33,14 +33,13 @@ logger = logging.getLogger(__name__)
 # Official leaderboard: https://huggingface.co/spaces/vidore/vidore-leaderboard
 VIDORE_DATASETS = {
     # === RECOMMENDED FOR QUICK TESTING (smaller, faster) ===
-    "docvqa": "vidore/docvqa_test_subsampled",      # ~500 queries, Document VQA
-    "infovqa": "vidore/infovqa_test_subsampled",    # ~500 queries, Infographics
+    "docvqa": "vidore/docvqa_test_subsampled",  # ~500 queries, Document VQA
+    "infovqa": "vidore/infovqa_test_subsampled",  # ~500 queries, Infographics
     "tabfquad": "vidore/tabfquad_test_subsampled",  # ~500 queries, Tables
-    
     # === FULL EVALUATION ===
-    "tatdqa": "vidore/tatdqa_test",                 # ~1500 queries, Financial tables
-    "arxivqa": "vidore/arxivqa_test_subsampled",    # ~500 queries, Scientific papers
-    "shift": "vidore/shiftproject_test",            # ~500 queries, Sustainability reports
+    "tatdqa": "vidore/tatdqa_test",  # ~1500 queries, Financial tables
+    "arxivqa": "vidore/arxivqa_test_subsampled",  # ~500 queries, Scientific papers
+    "shift": "vidore/shiftproject_test",  # ~500 queries, Sustainability reports
 }
 
 # Aliases for convenience
@@ -54,41 +53,45 @@ def load_dataset(dataset_name: str) -> Dict[str, Any]:
         from datasets import load_dataset
     except ImportError:
         raise ImportError("datasets library required. Install with: pip install datasets")
-    
+
     logger.info(f"Loading dataset: {dataset_name}")
-    
+
     # Load dataset
     ds = load_dataset(dataset_name, split="test")
-    
+
     # Extract queries and documents
     # ViDoRe format: each example has query, image, and relevant doc info
     queries = []
     documents = []
     qrels = {}  # query_id -> {doc_id: relevance}
-    
+
     for idx, example in enumerate(tqdm(ds, desc="Loading data")):
         query_id = f"q_{idx}"
         doc_id = f"d_{idx}"
-        
+
         # Get query text
         query_text = example.get("query", example.get("question", ""))
-        queries.append({
-            "id": query_id,
-            "text": query_text,
-        })
-        
+        queries.append(
+            {
+                "id": query_id,
+                "text": query_text,
+            }
+        )
+
         # Get document image
         image = example.get("image", example.get("page_image"))
-        documents.append({
-            "id": doc_id,
-            "image": image,
-        })
-        
+        documents.append(
+            {
+                "id": doc_id,
+                "image": image,
+            }
+        )
+
         # Relevance (self-document is relevant)
         qrels[query_id] = {doc_id: 1}
-    
+
     logger.info(f"Loaded {len(queries)} queries and {len(documents)} documents")
-    
+
     return {
         "queries": queries,
         "documents": documents,
@@ -104,30 +107,30 @@ def embed_documents(
 ) -> Dict[str, np.ndarray]:
     """
     Embed all documents.
-    
+
     Args:
         documents: List of {id, image} dicts
         embedder: VisualEmbedder instance
         batch_size: Batch size for embedding
         return_pooled: Also return tile-level pooled embeddings (for two-stage)
-    
+
     Returns:
         doc_embeddings dict, and optionally pooled_embeddings dict
     """
     from visual_rag.embedding.pooling import tile_level_mean_pooling
-    
+
     logger.info(f"Embedding {len(documents)} documents...")
-    
+
     images = [doc["image"] for doc in documents]
-    
+
     # Get embeddings with token info for proper pooling
     embeddings, token_infos = embedder.embed_images(
         images, batch_size=batch_size, return_token_info=True
     )
-    
+
     doc_embeddings = {}
     pooled_embeddings = {} if return_pooled else None
-    
+
     for doc, emb, token_info in zip(documents, embeddings, token_infos):
         if hasattr(emb, "numpy"):
             emb_np = emb.numpy()
@@ -135,18 +138,18 @@ def embed_documents(
             emb_np = emb.cpu().numpy()
         else:
             emb_np = np.array(emb)
-        
+
         doc_embeddings[doc["id"]] = emb_np.astype(np.float32)
-        
+
         # Compute tile-level pooling (NOVEL approach)
         if return_pooled:
             n_rows = token_info.get("n_rows", 4)
             n_cols = token_info.get("n_cols", 3)
             num_tiles = n_rows * n_cols + 1 if n_rows and n_cols else 13
-            
+
             pooled = tile_level_mean_pooling(emb_np, num_tiles, patches_per_tile=64)
             pooled_embeddings[doc["id"]] = pooled.astype(np.float32)
-    
+
     if return_pooled:
         return doc_embeddings, pooled_embeddings
     return doc_embeddings
@@ -158,7 +161,7 @@ def embed_queries(
 ) -> Dict[str, np.ndarray]:
     """Embed all queries."""
     logger.info(f"Embedding {len(queries)} queries...")
-    
+
     query_embeddings = {}
     for query in tqdm(queries, desc="Embedding queries"):
         emb = embedder.embed_query(query["text"])
@@ -167,7 +170,7 @@ def embed_queries(
         elif hasattr(emb, "cpu"):
             emb = emb.cpu().numpy()
         query_embeddings[query["id"]] = np.array(emb, dtype=np.float32)
-    
+
     return query_embeddings
 
 
@@ -176,10 +179,10 @@ def compute_maxsim(query_emb: np.ndarray, doc_emb: np.ndarray) -> float:
     # Normalize
     query_norm = query_emb / (np.linalg.norm(query_emb, axis=1, keepdims=True) + 1e-8)
     doc_norm = doc_emb / (np.linalg.norm(doc_emb, axis=1, keepdims=True) + 1e-8)
-    
+
     # Compute similarity matrix
     sim_matrix = np.dot(query_norm, doc_norm.T)
-    
+
     # MaxSim: max per query token, then sum
     max_sims = sim_matrix.max(axis=1)
     return float(max_sims.sum())
@@ -195,7 +198,7 @@ def search_exhaustive(
     for doc_id, doc_emb in doc_embeddings.items():
         score = compute_maxsim(query_emb, doc_emb)
         scores.append({"id": doc_id, "score": score})
-    
+
     # Sort by score
     scores.sort(key=lambda x: x["score"], reverse=True)
     return scores[:top_k]
@@ -210,11 +213,11 @@ def search_two_stage(
 ) -> List[Dict]:
     """
     Two-stage retrieval: tile-level pooled prefetch + MaxSim rerank.
-    
+
     Stage 1: Use tile-level pooled vectors for fast retrieval
              Each doc has [num_tiles, 128] pooled representation
              Compute MaxSim on pooled vectors (much faster)
-    
+
     Stage 2: Exact MaxSim reranking on top candidates
              Use full multi-vector embeddings for precision
     """
@@ -222,7 +225,7 @@ def search_two_stage(
     # Query pooled: mean across query tokens → [128]
     query_pooled = query_emb.mean(axis=0)
     query_pooled = query_pooled / (np.linalg.norm(query_pooled) + 1e-8)
-    
+
     stage1_scores = []
     for doc_id, doc_pooled in pooled_embeddings.items():
         # doc_pooled shape: [num_tiles, 128] from tile-level pooling
@@ -231,23 +234,25 @@ def search_two_stage(
         tile_sims = np.dot(doc_norm, query_pooled)
         score = float(tile_sims.max())  # Max tile similarity
         stage1_scores.append({"id": doc_id, "score": score})
-    
+
     stage1_scores.sort(key=lambda x: x["score"], reverse=True)
     candidates = stage1_scores[:prefetch_k]
-    
+
     # Stage 2: Exact MaxSim rerank on candidates
     reranked = []
     for cand in candidates:
         doc_id = cand["id"]
         doc_emb = doc_embeddings[doc_id]
         score = compute_maxsim(query_emb, doc_emb)
-        reranked.append({
-            "id": doc_id, 
-            "score": score, 
-            "stage1_score": cand["score"],
-            "stage1_rank": stage1_scores.index(cand) + 1,
-        })
-    
+        reranked.append(
+            {
+                "id": doc_id,
+                "score": score,
+                "stage1_score": cand["score"],
+                "stage1_rank": stage1_scores.index(cand) + 1,
+            }
+        )
+
     reranked.sort(key=lambda x: x["score"], reverse=True)
     return reranked[:top_k]
 
@@ -262,10 +267,10 @@ def compute_metrics(
     mrr_10 = []
     recall_5 = []
     recall_10 = []
-    
+
     for query_id, ranking in results.items():
         relevant = set(qrels.get(query_id, {}).keys())
-        
+
         # MRR@10
         rr = 0.0
         for i, doc in enumerate(ranking[:10]):
@@ -273,32 +278,30 @@ def compute_metrics(
                 rr = 1.0 / (i + 1)
                 break
         mrr_10.append(rr)
-        
+
         # Recall@5, Recall@10
         retrieved_5 = set(d["id"] for d in ranking[:5])
         retrieved_10 = set(d["id"] for d in ranking[:10])
-        
+
         if relevant:
             recall_5.append(len(retrieved_5 & relevant) / len(relevant))
             recall_10.append(len(retrieved_10 & relevant) / len(relevant))
-        
+
         # NDCG@5, NDCG@10
-        dcg_5 = sum(
-            1.0 / np.log2(i + 2) for i, d in enumerate(ranking[:5]) if d["id"] in relevant
-        )
+        dcg_5 = sum(1.0 / np.log2(i + 2) for i, d in enumerate(ranking[:5]) if d["id"] in relevant)
         dcg_10 = sum(
             1.0 / np.log2(i + 2) for i, d in enumerate(ranking[:10]) if d["id"] in relevant
         )
-        
+
         # Ideal DCG
         k_rel = min(len(relevant), 5)
         idcg_5 = sum(1.0 / np.log2(i + 2) for i in range(k_rel))
         k_rel = min(len(relevant), 10)
         idcg_10 = sum(1.0 / np.log2(i + 2) for i in range(k_rel))
-        
+
         ndcg_5.append(dcg_5 / idcg_5 if idcg_5 > 0 else 0.0)
         ndcg_10.append(dcg_10 / idcg_10 if idcg_10 > 0 else 0.0)
-    
+
     return {
         "ndcg@5": float(np.mean(ndcg_5)),
         "ndcg@10": float(np.mean(ndcg_10)),
@@ -318,87 +321,90 @@ def run_evaluation(
 ) -> Dict[str, Any]:
     """Run full evaluation on a dataset."""
     from visual_rag.embedding import VisualEmbedder
-    
-    logger.info(f"=" * 60)
+
+    logger.info("=" * 60)
     logger.info(f"Evaluating: {dataset_name}")
     logger.info(f"Model: {model_name}")
     logger.info(f"Two-stage: {two_stage}")
-    logger.info(f"=" * 60)
-    
+    logger.info("=" * 60)
+
     # Load dataset
     data = load_dataset(dataset_name)
-    
+
     # Initialize embedder
     embedder = VisualEmbedder(model_name=model_name)
-    
+
     # Embed documents (with tile-level pooling if two-stage)
     start_time = time.time()
     if two_stage:
         doc_embeddings, pooled_embeddings = embed_documents(
             data["documents"], embedder, return_pooled=True
         )
-        logger.info(f"Using tile-level pooling for two-stage retrieval")
+        logger.info("Using tile-level pooling for two-stage retrieval")
     else:
         doc_embeddings = embed_documents(data["documents"], embedder)
         pooled_embeddings = None
     embed_time = time.time() - start_time
     logger.info(f"Document embedding time: {embed_time:.2f}s")
-    
+
     # Embed queries
     query_embeddings = embed_queries(data["queries"], embedder)
-    
+
     # Run search
     logger.info("Running search...")
     results = {}
     search_times = []
-    
+
     for query in tqdm(data["queries"], desc="Searching"):
         query_id = query["id"]
         query_emb = query_embeddings[query_id]
-        
+
         start = time.time()
         if two_stage:
             ranking = search_two_stage(
-                query_emb, doc_embeddings, pooled_embeddings,
-                prefetch_k=prefetch_k, top_k=top_k
+                query_emb, doc_embeddings, pooled_embeddings, prefetch_k=prefetch_k, top_k=top_k
             )
         else:
             ranking = search_exhaustive(query_emb, doc_embeddings, top_k=top_k)
         search_times.append(time.time() - start)
-        
+
         results[query_id] = ranking
-    
+
     avg_search_time = np.mean(search_times)
     logger.info(f"Average search time: {avg_search_time * 1000:.2f}ms")
-    
+
     # Compute metrics
     metrics = compute_metrics(results, data["qrels"])
     metrics["avg_search_time_ms"] = avg_search_time * 1000
     metrics["embed_time_s"] = embed_time
-    
-    logger.info(f"\nResults:")
+
+    logger.info("\nResults:")
     for k, v in metrics.items():
         logger.info(f"  {k}: {v:.4f}")
-    
+
     # Save results
     if output_dir:
         output_path = Path(output_dir)
         output_path.mkdir(parents=True, exist_ok=True)
-        
+
         dataset_short = dataset_name.split("/")[-1]
         suffix = "_twostage" if two_stage else ""
         result_file = output_path / f"{dataset_short}{suffix}.json"
-        
+
         with open(result_file, "w") as f:
-            json.dump({
-                "dataset": dataset_name,
-                "model": model_name,
-                "two_stage": two_stage,
-                "metrics": metrics,
-            }, f, indent=2)
-        
+            json.dump(
+                {
+                    "dataset": dataset_name,
+                    "model": model_name,
+                    "two_stage": two_stage,
+                    "metrics": metrics,
+                },
+                f,
+                indent=2,
+            )
+
         logger.info(f"Saved results to: {result_file}")
-    
+
     return metrics
 
 
@@ -413,57 +419,53 @@ Available datasets:
 Examples:
   # Quick test on DocVQA
   python run_vidore.py --dataset docvqa
-  
+
   # Quick test with two-stage (your novel approach)
   python run_vidore.py --dataset docvqa --two-stage
-  
+
   # Run on recommended quick datasets
   python run_vidore.py --quick
-  
+
   # Full evaluation on all datasets
   python run_vidore.py --all
-  
+
   # Compare exhaustive vs two-stage
   python run_vidore.py --dataset docvqa
   python run_vidore.py --dataset docvqa --two-stage
   python analyze_results.py --results results/ --compare
-"""
-    )
-    parser.add_argument(
-        "--dataset", type=str, choices=list(VIDORE_DATASETS.keys()),
-        help=f"Dataset to evaluate: {', '.join(VIDORE_DATASETS.keys())}"
-    )
-    parser.add_argument(
-        "--quick", action="store_true",
-        help=f"Run on quick datasets: {QUICK_DATASETS}"
+""",
     )
     parser.add_argument(
-        "--all", action="store_true",
-        help="Evaluate on all ViDoRe datasets"
+        "--dataset",
+        type=str,
+        choices=list(VIDORE_DATASETS.keys()),
+        help=f"Dataset to evaluate: {', '.join(VIDORE_DATASETS.keys())}",
     )
     parser.add_argument(
-        "--model", type=str, default="vidore/colSmol-500M",
-        help="Model: vidore/colSmol-500M (default), vidore/colpali-v1.3, vidore/colqwen2-v1.0"
+        "--quick", action="store_true", help=f"Run on quick datasets: {QUICK_DATASETS}"
     )
+    parser.add_argument("--all", action="store_true", help="Evaluate on all ViDoRe datasets")
     parser.add_argument(
-        "--two-stage", action="store_true",
-        help="Use two-stage retrieval (tile-level pooled prefetch + MaxSim rerank)"
+        "--model",
+        type=str,
+        default="vidore/colSmol-500M",
+        help="Model: vidore/colSmol-500M (default), vidore/colpali-v1.3, vidore/colqwen2-v1.0",
     )
     parser.add_argument(
-        "--prefetch-k", type=int, default=100,
-        help="Stage 1 candidates (default: 100)"
+        "--two-stage",
+        action="store_true",
+        help="Use two-stage retrieval (tile-level pooled prefetch + MaxSim rerank)",
     )
     parser.add_argument(
-        "--top-k", type=int, default=10,
-        help="Final results (default: 10)"
+        "--prefetch-k", type=int, default=100, help="Stage 1 candidates (default: 100)"
     )
+    parser.add_argument("--top-k", type=int, default=10, help="Final results (default: 10)")
     parser.add_argument(
-        "--output-dir", type=str, default="results",
-        help="Output directory (default: results)"
+        "--output-dir", type=str, default="results", help="Output directory (default: results)"
     )
-    
+
     args = parser.parse_args()
-    
+
     # Determine which datasets to run
     if args.all:
         dataset_keys = ALL_DATASETS
@@ -473,11 +475,11 @@ Examples:
         dataset_keys = [args.dataset]
     else:
         parser.error("Specify --dataset, --quick, or --all")
-    
+
     # Convert keys to full HuggingFace paths
     datasets = [VIDORE_DATASETS[k] for k in dataset_keys]
     logger.info(f"Running on {len(datasets)} dataset(s): {dataset_keys}")
-    
+
     all_results = {}
     for dataset in datasets:
         try:
@@ -493,21 +495,19 @@ Examples:
         except Exception as e:
             logger.error(f"Failed on {dataset}: {e}")
             continue
-    
+
     # Summary
     if len(all_results) > 1:
         logger.info("\n" + "=" * 60)
         logger.info("SUMMARY")
         logger.info("=" * 60)
-        
+
         avg_ndcg10 = np.mean([m["ndcg@10"] for m in all_results.values()])
         avg_mrr10 = np.mean([m["mrr@10"] for m in all_results.values()])
-        
+
         logger.info(f"Average NDCG@10: {avg_ndcg10:.4f}")
         logger.info(f"Average MRR@10: {avg_mrr10:.4f}")
 
 
 if __name__ == "__main__":
     main()
-
-

benchmarks/vidore_beir_qdrant/run_qdrant_beir.py

@@ -4,13 +4,14 @@ import os
 import sys
 import tempfile
 import time
+import warnings
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple
 
 import numpy as np
 
 from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
-from benchmarks.vidore_tatdqa_test.metrics import ndcg_at_k, mrr_at_k, recall_at_k
+from benchmarks.vidore_tatdqa_test.metrics import mrr_at_k, ndcg_at_k, recall_at_k
 from visual_rag import VisualEmbedder
 from visual_rag.indexing.cloudinary_uploader import CloudinaryUploader
 from visual_rag.indexing.qdrant_indexer import QdrantIndexer
@@ -83,15 +84,20 @@ def _parse_payload_indexes(values: List[str]) -> List[Dict[str, str]]:
     return indexes
 
 
-def _union_point_id(*, dataset_name: str, source_doc_id: str, union_namespace: Optional[str]) -> str:
+def _union_point_id(
+    *, dataset_name: str, source_doc_id: str, union_namespace: Optional[str]
+) -> str:
     ns = f"{union_namespace}::{dataset_name}" if union_namespace else dataset_name
     return _stable_uuid(f"{ns}::{source_doc_id}")
 
 
-def _filter_qrels(qrels: Dict[str, Dict[str, int]], query_ids: List[str]) -> Dict[str, Dict[str, int]]:
+def _filter_qrels(
+    qrels: Dict[str, Dict[str, int]], query_ids: List[str]
+) -> Dict[str, Dict[str, int]]:
     keep = set(query_ids)
     return {qid: rels for qid, rels in qrels.items() if qid in keep}
 
+
 def _failed_log_path(*, collection_name: str, dataset_name: str) -> Path:
     dir_name = _safe_filename(collection_name)
     return Path("results") / dir_name / f"index_failures__{_safe_filename(dataset_name)}.jsonl"
@@ -130,7 +136,7 @@ def _default_output_filename(*, args, datasets: List[str]) -> str:
     if str(args.mode) == "three_stage":
         parts.append("tokens_vs_global")
         parts.append(f"s1k{int(args.stage1_k)}")
-        parts.append("tokens_vs_experimental")
+        parts.append("tokens_vs_experimental_pooling")
         parts.append(f"s2k{int(args.stage2_k)}")
     parts.extend([topk_tag, scope_tag, ds_tag])
 
@@ -207,7 +213,9 @@ def _load_failed_union_ids(
     return out
 
 
-def _remove_failed_from_qrels(qrels: Dict[str, Dict[str, int]], failed_ids: set) -> Tuple[Dict[str, Dict[str, int]], int]:
+def _remove_failed_from_qrels(
+    qrels: Dict[str, Dict[str, int]], failed_ids: set
+) -> Tuple[Dict[str, Dict[str, int]], int]:
     removed = 0
     if not failed_ids:
         return qrels, 0
@@ -223,6 +231,45 @@ def _remove_failed_from_qrels(qrels: Dict[str, Dict[str, int]], failed_ids: set)
     return out, removed
 
 
+def _filter_failed_ids_to_missing(
+    *,
+    qdrant_client,
+    collection_name: str,
+    failed_ids: set,
+    timeout: int,
+    batch_size: int = 128,
+) -> set:
+    """
+    Failure logs are append-only and can contain historical IDs that may have been
+    successfully retried later. To avoid poisoning evaluation, keep only IDs that
+    are *still missing* in Qdrant.
+    """
+    failed_ids = set(str(x) for x in (failed_ids or set()) if x)
+    if not failed_ids:
+        return set()
+
+    missing = set()
+    ids_list = list(failed_ids)
+    for i in range(0, len(ids_list), int(batch_size)):
+        chunk = ids_list[i : i + int(batch_size)]
+        try:
+            recs = qdrant_client.retrieve(
+                collection_name=str(collection_name),
+                ids=chunk,
+                with_payload=False,
+                with_vectors=False,
+                timeout=int(timeout),
+            )
+            present = set(str(r.id) for r in (recs or []))
+            for cid in chunk:
+                if str(cid) not in present:
+                    missing.add(str(cid))
+        except Exception:
+            # If retrieve fails (e.g. transient network), be conservative: treat as missing.
+            missing.update(str(cid) for cid in chunk)
+    return missing
+
+
 def _evaluate(
     *,
     queries,
@@ -283,11 +330,25 @@ def _evaluate(
     retrieve_k = max(100, top_k)
 
     query_texts = [q.text for q in queries]
-    query_embeddings = embedder.embed_queries(
-        query_texts,
-        batch_size=getattr(embedder, "batch_size", None),
-        show_progress=False,
-    )
+    embed_started_at = time.time()
+    print(f"📝 Embedding {len(query_texts)} queries…")
+    sys.stdout.flush()
+    # Always try to show a progress bar during query embedding.
+    # If the installed VisualEmbedder version doesn't support show_progress, fall back gracefully.
+    try:
+        query_embeddings = embedder.embed_queries(
+            query_texts,
+            batch_size=getattr(embedder, "batch_size", None),
+            show_progress=True,
+        )
+    except TypeError:
+        query_embeddings = embedder.embed_queries(
+            query_texts,
+            batch_size=getattr(embedder, "batch_size", None),
+        )
+    embed_s = float(max(time.time() - embed_started_at, 0.0))
+    print(f"✅ Embedded queries in {embed_s:.2f}s")
+    sys.stdout.flush()
 
     iterator = queries
     try:
@@ -304,7 +365,10 @@ def _evaluate(
         except ImportError:
             torch = None
         if torch is not None and isinstance(qemb, torch.Tensor):
-            qemb_np = qemb.detach().cpu().numpy()
+            # Keep evaluation stable across dtypes/devices:
+            # - numpy doesn't support bfloat16
+            # - float16 queries can cause large quality drops on some backends
+            qemb_np = qemb.detach().float().cpu().numpy()
         else:
             qemb_np = qemb.numpy()
 
@@ -361,6 +425,41 @@ def _evaluate(
     }
 
 
+def _detect_collection_vector_dtype(*, client, collection_name: str) -> Optional[str]:
+    """
+    Best-effort detection of the stored vector datatype for a Qdrant collection.
+
+    Returns:
+        "float16", "float32", or None if unavailable.
+    """
+    try:
+        info = client.get_collection(str(collection_name))
+    except Exception:
+        return None
+
+    try:
+        vectors = info.config.params.vectors or {}
+    except Exception:
+        vectors = {}
+
+    vp = None
+    if isinstance(vectors, dict):
+        vp = vectors.get("initial") or (next(iter(vectors.values())) if vectors else None)
+    if vp is None:
+        return None
+
+    dt = getattr(vp, "datatype", None)
+    if dt is None:
+        return None
+
+    s = str(dt).lower()
+    if "float16" in s:
+        return "float16"
+    if "float32" in s:
+        return "float32"
+    return None
+
+
 def _write_json_atomic(path: Path, data: Dict[str, Any]) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
     fd, tmp_path = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent))
@@ -415,12 +514,15 @@ def _index_beir_corpus(
     crop_empty_remove_page_number: bool,
     crop_empty_preserve_border_px: int,
     crop_empty_uniform_std_threshold: float,
+    max_mean_pool_vectors: Optional[int],
     no_cloudinary: bool,
     cloudinary_folder: str,
     retry_failures: bool,
     only_failures: bool,
 ) -> None:
-    qdrant_url = os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    qdrant_url = (
+        os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL")
+    )
     if not qdrant_url:
         raise ValueError("QDRANT_URL not set")
     qdrant_api_key = (
@@ -453,7 +555,9 @@ def _index_beir_corpus(
             cloudinary_uploader = None
 
     failure_log = _failed_log_path(collection_name=collection_name, dataset_name=dataset_name)
-    failed_ids = _load_failed_union_ids(failure_log, dataset_name=dataset_name, union_namespace=union_namespace)
+    failed_ids = _load_failed_union_ids(
+        failure_log, dataset_name=dataset_name, union_namespace=union_namespace
+    )
     previously_failed_ids = set(failed_ids)
 
     existing_ids = set()
@@ -520,6 +624,7 @@ def _index_beir_corpus(
         out = img.copy()
         out.thumbnail((1024, 1024), Image.BICUBIC)
         return out
+
     uploaded_docs = 0
     skipped_docs = 0
     start_time = time.time()
@@ -536,14 +641,24 @@ def _index_beir_corpus(
         pass
 
     import threading
-    from concurrent.futures import ThreadPoolExecutor, wait as futures_wait, FIRST_EXCEPTION
+    from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor
+    from concurrent.futures import wait as futures_wait
 
     stop_event = threading.Event()
-    executor = ThreadPoolExecutor(max_workers=int(upload_workers)) if upload_workers and upload_workers > 0 else None
+    executor = (
+        ThreadPoolExecutor(max_workers=int(upload_workers))
+        if upload_workers and upload_workers > 0
+        else None
+    )
     futures = []
 
     def _upload(points: List[Dict[str, Any]]) -> int:
-        uploaded = int(indexer.upload_batch(points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event) or 0)
+        uploaded = int(
+            indexer.upload_batch(
+                points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event
+            )
+            or 0
+        )
         if uploaded <= 0 and points:
             for p in points:
                 pid = str(p.get("id") or "")
@@ -554,7 +669,9 @@ def _index_beir_corpus(
                             "dataset": dataset_name,
                             "collection": collection_name,
                             "model": model_name,
-                            "source_doc_id": str((p.get("metadata") or {}).get("source_doc_id") or ""),
+                            "source_doc_id": str(
+                                (p.get("metadata") or {}).get("source_doc_id") or ""
+                            ),
                             "doc_id": str((p.get("metadata") or {}).get("doc_id") or ""),
                             "union_doc_id": pid,
                             "error": "Qdrant upsert failed (all retries exhausted)",
@@ -632,7 +749,8 @@ def _index_beir_corpus(
                 continue
 
             if crop_empty:
-                from visual_rag.preprocessing.crop_empty import CropEmptyConfig, crop_empty as _crop_empty
+                from visual_rag.preprocessing.crop_empty import CropEmptyConfig
+                from visual_rag.preprocessing.crop_empty import crop_empty as _crop_empty
 
                 crop_cfg = CropEmptyConfig(
                     percentage_to_remove=float(crop_empty_percentage_to_remove),
@@ -660,7 +778,7 @@ def _index_beir_corpus(
                     return_token_info=True,
                     show_progress=False,
                 )
-            except Exception as e:
+            except Exception:
                 # Retry per-doc to isolate flaky backend / corrupted sample issues.
                 embeddings = []
                 token_infos = []
@@ -675,7 +793,9 @@ def _index_beir_corpus(
                         embeddings.append(e1[0])
                         token_infos.append(t1[0])
                     except Exception as e_single:
-                        source_doc_id_i = str((doc_i.payload or {}).get("source_doc_id") or doc_i.doc_id)
+                        source_doc_id_i = str(
+                            (doc_i.payload or {}).get("source_doc_id") or doc_i.doc_id
+                        )
                         union_doc_id_i = _union_point_id(
                             dataset_name=dataset_name,
                             source_doc_id=source_doc_id_i,
@@ -704,63 +824,107 @@ def _index_beir_corpus(
             for doc, emb, token_info, crop_meta, original_img, embed_img in zip(
                 batch, embeddings, token_infos, crop_metas, original_images, images
             ):
+                source_doc_id = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
+                union_doc_id = _union_point_id(
+                    dataset_name=dataset_name,
+                    source_doc_id=source_doc_id,
+                    union_namespace=union_namespace,
+                )
+
                 try:
-                    emb_np = emb.cpu().float().numpy() if hasattr(emb, "cpu") else np.array(emb, dtype=np.float32)
-                    visual_indices = token_info.get("visual_token_indices") or list(range(emb_np.shape[0]))
+                    emb_np = (
+                        emb.cpu().float().numpy()
+                        if hasattr(emb, "cpu")
+                        else np.array(emb, dtype=np.float32)
+                    )
+                    visual_indices = token_info.get("visual_token_indices") or list(
+                        range(emb_np.shape[0])
+                    )
                     visual_embedding = emb_np[visual_indices].astype(np.float32)
-                    tile_pooled = embedder.mean_pool_visual_embedding(visual_embedding, token_info, target_vectors=32)
+                    tile_pooled = embedder.mean_pool_visual_embedding(
+                        visual_embedding, token_info, target_vectors=max_mean_pool_vectors
+                    )
                     experimental_pooled = embedder.experimental_pool_visual_embedding(
-                        visual_embedding, token_info, target_vectors=32, mean_pool=tile_pooled
+                        visual_embedding,
+                        token_info,
+                        target_vectors=max_mean_pool_vectors,
+                        mean_pool=tile_pooled,
                     )
                     global_pooled = embedder.global_pool_from_mean_pool(tile_pooled)
+
+                    # Log whenever ColQwen2.5 adaptive mean pooling actually downsamples rows.
+                    model_lower = (model_name or "").lower()
+                    is_colqwen25 = "colqwen2.5" in model_lower or "colqwen2_5" in model_lower
+                    if is_colqwen25:
+                        grid_h_eff = (token_info or {}).get("grid_h_eff")
+                        if grid_h_eff is not None:
+                            try:
+                                h_eff = int(grid_h_eff)
+                                out_rows = int(getattr(tile_pooled, "shape", [0])[0])
+                            except Exception:
+                                h_eff = 0
+                                out_rows = 0
+                            if h_eff > 0 and out_rows > 0 and out_rows < h_eff:
+                                msg = (
+                                    "Downsampled ColQwen mean-pool rows for "
+                                    f"union_doc_id={union_doc_id} (source_doc_id={source_doc_id}): "
+                                    f"grid_h_eff={h_eff} -> {out_rows} "
+                                    f"(--max-mean-pool-vectors={max_mean_pool_vectors})"
+                                )
+                                if pbar is not None:
+                                    try:
+                                        pbar.write(msg)
+                                    except Exception:
+                                        print(msg)
+                                else:
+                                    print(msg)
                 except Exception as e_single:
-                    source_doc_id_i = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
-                    union_doc_id_i = _union_point_id(
-                        dataset_name=dataset_name,
-                        source_doc_id=source_doc_id_i,
-                        union_namespace=union_namespace,
-                    )
-                    if str(union_doc_id_i) not in failed_ids:
+                    if str(union_doc_id) not in failed_ids:
                         _append_jsonl(
                             failure_log,
                             {
                                 "dataset": dataset_name,
                                 "collection": collection_name,
                                 "model": model_name,
-                                "source_doc_id": str(source_doc_id_i),
+                                "source_doc_id": str(source_doc_id),
                                 "doc_id": str(getattr(doc, "doc_id", "")),
-                                "union_doc_id": str(union_doc_id_i),
+                                "union_doc_id": str(union_doc_id),
                                 "error": str(e_single),
                             },
                         )
-                        failed_ids.add(str(union_doc_id_i))
-                    existing_ids.add(str(union_doc_id_i))
+                        failed_ids.add(str(union_doc_id))
+                    existing_ids.add(str(union_doc_id))
                     skipped_docs += 1
                     continue
 
                 num_tiles = int(tile_pooled.shape[0])
-                patches_per_tile = int(visual_embedding.shape[0] // max(num_tiles, 1)) if num_tiles else 0
-
-                source_doc_id = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
-                union_doc_id = _union_point_id(
-                    dataset_name=dataset_name,
-                    source_doc_id=source_doc_id,
-                    union_namespace=union_namespace,
+                patches_per_tile = (
+                    int(visual_embedding.shape[0] // max(num_tiles, 1)) if num_tiles else 0
                 )
 
                 resized_img = _resized_for_display(embed_img) or embed_img
                 original_url = ""
                 cropped_url = ""
                 resized_url = ""
-                if cloudinary_uploader is not None and original_img is not None and resized_img is not None:
+                if (
+                    cloudinary_uploader is not None
+                    and original_img is not None
+                    and resized_img is not None
+                ):
                     base_public_id = _safe_public_id(f"{dataset_name}__{union_doc_id}")
                     try:
                         if crop_empty:
-                            o_url, c_url, r_url = cloudinary_uploader.upload_original_cropped_and_resized(
-                                original_img,
-                                embed_img if embed_img is not None and embed_img is not original_img else None,
-                                resized_img,
-                                base_public_id,
+                            o_url, c_url, r_url = (
+                                cloudinary_uploader.upload_original_cropped_and_resized(
+                                    original_img,
+                                    (
+                                        embed_img
+                                        if embed_img is not None and embed_img is not original_img
+                                        else None
+                                    ),
+                                    resized_img,
+                                    base_public_id,
+                                )
                             )
                             original_url = o_url or ""
                             cropped_url = c_url or ""
@@ -782,12 +946,18 @@ def _index_beir_corpus(
                     "union_doc_id": union_doc_id,
                     "page": resized_url or original_url or "",
                     "original_url": original_url,
-                    "cropped_url": cropped_url,
+                    "cropped_url": cropped_url if crop_empty else "",
                     "resized_url": resized_url,
                     "original_width": int(original_img.width) if original_img is not None else None,
-                    "original_height": int(original_img.height) if original_img is not None else None,
-                    "cropped_width": int(embed_img.width) if embed_img is not None else None,
-                    "cropped_height": int(embed_img.height) if embed_img is not None else None,
+                    "original_height": (
+                        int(original_img.height) if original_img is not None else None
+                    ),
+                    "cropped_width": (
+                        int(embed_img.width) if (crop_empty and embed_img is not None) else None
+                    ),
+                    "cropped_height": (
+                        int(embed_img.height) if (crop_empty and embed_img is not None) else None
+                    ),
                     "resized_width": int(resized_img.width) if resized_img is not None else None,
                     "resized_height": int(resized_img.height) if resized_img is not None else None,
                     "num_tiles": int(num_tiles),
@@ -795,15 +965,25 @@ def _index_beir_corpus(
                     "torch_dtype": _torch_dtype_to_str(embedder.torch_dtype),
                     "model_name": model_name,
                     "crop_empty_enabled": bool(crop_empty),
-                    "crop_empty_crop_box": (crop_meta or {}).get("crop_box") if crop_empty else None,
-                    "crop_empty_remove_page_number": bool(crop_empty_remove_page_number) if crop_empty else None,
-                    "crop_empty_percentage_to_remove": float(crop_empty_percentage_to_remove) if crop_empty else None,
+                    "crop_empty_crop_box": (
+                        (crop_meta or {}).get("crop_box") if crop_empty else None
+                    ),
+                    "crop_empty_remove_page_number": (
+                        bool(crop_empty_remove_page_number) if crop_empty else None
+                    ),
+                    "crop_empty_percentage_to_remove": (
+                        float(crop_empty_percentage_to_remove) if crop_empty else None
+                    ),
                     "index_recovery_previously_failed": bool(union_doc_id in previously_failed_ids),
                     "index_recovery_mode": (
-                        "only_failures" if bool(only_failures) else ("retry_failures" if bool(retry_failures) else None)
+                        "only_failures"
+                        if bool(only_failures)
+                        else ("retry_failures" if bool(retry_failures) else None)
                     ),
                     "index_recovery_pooling_inferred_tiles": bool(
-                        (token_info or {}).get("num_tiles") is None and (token_info or {}).get("n_rows") is None and (token_info or {}).get("n_cols") is None
+                        (token_info or {}).get("num_tiles") is None
+                        and (token_info or {}).get("n_rows") is None
+                        and (token_info or {}).get("n_cols") is None
                     ),
                     "index_recovery_num_visual_tokens": int(visual_embedding.shape[0]),
                     **(doc.payload or {}),
@@ -894,8 +1074,8 @@ def main() -> None:
     parser.add_argument(
         "--qdrant-vector-dtype",
         type=str,
-        default="float16",
-        choices=["float16", "float32"],
+        default="auto",
+        choices=["auto", "float16", "float32"],
     )
     grpc_group = parser.add_mutually_exclusive_group()
     grpc_group.add_argument("--prefer-grpc", dest="prefer_grpc", action="store_true", default=True)
@@ -910,10 +1090,14 @@ def main() -> None:
     parser.add_argument("--upsert-wait", action="store_true")
     parser.add_argument("--max-corpus-docs", type=int, default=0)
     parser.add_argument("--sample-corpus-docs", type=int, default=0)
-    parser.add_argument("--sample-corpus-strategy", type=str, default="first", choices=["first", "random"])
+    parser.add_argument(
+        "--sample-corpus-strategy", type=str, default="first", choices=["first", "random"]
+    )
     parser.add_argument("--sample-seed", type=int, default=0)
     parser.add_argument("--sample-queries", type=int, default=0)
-    parser.add_argument("--sample-query-strategy", type=str, default="first", choices=["first", "random"])
+    parser.add_argument(
+        "--sample-query-strategy", type=str, default="first", choices=["first", "random"]
+    )
     parser.add_argument("--sample-query-seed", type=int, default=0)
     parser.add_argument("--index-from-queries", action="store_true", default=False)
     parser.add_argument("--resume", action="store_true", default=False)
@@ -925,14 +1109,35 @@ def main() -> None:
     parser.add_argument("--crop-empty-remove-page-number", action="store_true", default=False)
     parser.add_argument("--crop-empty-preserve-border-px", type=int, default=1)
     parser.add_argument("--crop-empty-uniform-std-threshold", type=float, default=0.0)
+    parser.add_argument(
+        "--max-mean-pool-vectors",
+        type=int,
+        default=None,
+        help=(
+            "Cap ColQwen2.5 adaptive row-mean pooling to at most this many vectors. "
+            "If omitted (default), no cap is applied (use all effective rows). "
+            "If <= 0, treated as no cap."
+        ),
+    )
     payload_group = parser.add_mutually_exclusive_group()
     payload_group.add_argument("--index-common-metadata", action="store_true", default=True)
-    payload_group.add_argument("--no-index-common-metadata", dest="index_common_metadata", action="store_false")
+    payload_group.add_argument(
+        "--no-index-common-metadata", dest="index_common_metadata", action="store_false"
+    )
     parser.add_argument("--payload-index", action="append", default=[])
-    parser.add_argument(
+    cloud_group = parser.add_mutually_exclusive_group()
+    cloud_group.add_argument(
+        "--cloudinary",
+        dest="no_cloudinary",
+        action="store_false",
+        default=True,
+        help="Enable Cloudinary uploads during indexing (default: disabled).",
+    )
+    cloud_group.add_argument(
         "--no-cloudinary",
+        dest="no_cloudinary",
         action="store_true",
-        help="Disable Cloudinary uploads during indexing (default: enabled).",
+        help="Disable Cloudinary uploads during indexing (default).",
     )
     parser.add_argument(
         "--cloudinary-folder",
@@ -953,8 +1158,18 @@ def main() -> None:
         help="Index only documents listed in index_failures__<collection>__<dataset>.jsonl.",
     )
 
-    parser.add_argument("--top-k", type=int, default=100, help="Retrieve top-k results (default: 100 to calculate metrics at all cutoffs)")
-    parser.add_argument("--prefetch-k", type=int, default=200, help="Prefetch candidates for two-stage (default: 200)")
+    parser.add_argument(
+        "--top-k",
+        type=int,
+        default=100,
+        help="Retrieve top-k results (default: 100 to calculate metrics at all cutoffs)",
+    )
+    parser.add_argument(
+        "--prefetch-k",
+        type=int,
+        default=200,
+        help="Prefetch candidates for two-stage (default: 200)",
+    )
     parser.add_argument(
         "--no-eval",
         action="store_true",
@@ -970,21 +1185,41 @@ def main() -> None:
     parser.add_argument(
         "--stage1-mode",
         type=str,
-        default="tokens_vs_tiles",
+        default="tokens_vs_standard_pooling",
         choices=[
+            # New naming (preferred)
+            "pooled_query_vs_standard_pooling",
+            "tokens_vs_standard_pooling",
+            "pooled_query_vs_experimental_pooling",
+            "tokens_vs_experimental_pooling",
+            "pooled_query_vs_global",
+            # Backwards-compatible aliases (deprecated)
             "pooled_query_vs_tiles",
             "tokens_vs_tiles",
             "pooled_query_vs_experimental",
             "tokens_vs_experimental",
-            "pooled_query_vs_global",
         ],
+        help=(
+            "Two-stage stage1 prefetch mode. "
+            "standard_pooling uses Qdrant named vector 'mean_pooling'. "
+            "experimental_pooling uses Qdrant named vector 'experimental_pooling'. "
+            "global uses Qdrant named vector 'global_pooling'."
+        ),
+    )
+    parser.add_argument(
+        "--stage1-k", type=int, default=1000, help="Three-stage stage1 top_k (default: 1000)"
+    )
+    parser.add_argument(
+        "--stage2-k", type=int, default=300, help="Three-stage stage2 top_k (default: 300)"
     )
-    parser.add_argument("--stage1-k", type=int, default=1000, help="Three-stage stage1 top_k (default: 1000)")
-    parser.add_argument("--stage2-k", type=int, default=300, help="Three-stage stage2 top_k (default: 300)")
     parser.add_argument("--max-queries", type=int, default=0)
     drop_group = parser.add_mutually_exclusive_group()
-    drop_group.add_argument("--drop-empty-queries", dest="drop_empty_queries", action="store_true", default=True)
-    drop_group.add_argument("--no-drop-empty-queries", dest="drop_empty_queries", action="store_false")
+    drop_group.add_argument(
+        "--drop-empty-queries", dest="drop_empty_queries", action="store_true", default=True
+    )
+    drop_group.add_argument(
+        "--no-drop-empty-queries", dest="drop_empty_queries", action="store_false"
+    )
     parser.add_argument(
         "--evaluation-scope",
         type=str,
@@ -1008,25 +1243,60 @@ def main() -> None:
         help="Stop the run immediately on the first dataset evaluation failure.",
     )
     parser.add_argument("--output", type=str, default="auto")
+    parser.add_argument(
+        "--ensure-in-ram",
+        dest="ensure_in_ram",
+        action="store_true",
+        default=False,
+        help="Best-effort: patch collection config so vectors/indexes are stored in RAM (on_disk=false).",
+    )
 
     args = parser.parse_args()
 
+    # Backwards-compatible stage1_mode mapping (deprecated names)
+    stage1_map = {
+        "pooled_query_vs_tiles": "pooled_query_vs_standard_pooling",
+        "tokens_vs_tiles": "tokens_vs_standard_pooling",
+        "pooled_query_vs_experimental": "pooled_query_vs_experimental_pooling",
+        "tokens_vs_experimental": "tokens_vs_experimental_pooling",
+    }
+    if str(args.stage1_mode) in stage1_map:
+        old = str(args.stage1_mode)
+        new = stage1_map[old]
+        warnings.warn(
+            f"--stage1-mode {old} is deprecated; use {new} instead.",
+            category=DeprecationWarning,
+            stacklevel=2,
+        )
+        args.stage1_mode = new
+
     _maybe_load_dotenv()
 
     if args.recreate:
         args.index = True
 
-    if args.sample_corpus_docs and int(args.sample_corpus_docs) > 0 and args.max_corpus_docs and int(args.max_corpus_docs) > 0:
+    if (
+        args.sample_corpus_docs
+        and int(args.sample_corpus_docs) > 0
+        and args.max_corpus_docs
+        and int(args.max_corpus_docs) > 0
+    ):
         raise ValueError("Use only one of --sample-corpus-docs or --max-corpus-docs (not both).")
     if args.sample_queries and int(args.sample_queries) > 0 and args.index_from_queries:
-        if (args.sample_corpus_docs and int(args.sample_corpus_docs) > 0) or (args.max_corpus_docs and int(args.max_corpus_docs) > 0):
-            raise ValueError("Use --index-from-queries with --sample-queries only (do not combine with corpus sampling).")
+        if (args.sample_corpus_docs and int(args.sample_corpus_docs) > 0) or (
+            args.max_corpus_docs and int(args.max_corpus_docs) > 0
+        ):
+            raise ValueError(
+                "Use --index-from-queries with --sample-queries only (do not combine with corpus sampling)."
+            )
 
     if args.upsert_wait:
         print("Qdrant upserts wait for completion (wait=True).")
     else:
         print("Qdrant upserts are async (wait=False).")
-    print(f"Qdrant request timeout: {int(args.qdrant_timeout)}s, retries: {int(args.qdrant_retries)}.")
+    print(
+        f"Qdrant request timeout: {int(args.qdrant_timeout)}s, retries: {int(args.qdrant_retries)}."
+    )
 
     datasets: List[str] = []
     if args.datasets:
@@ -1044,7 +1314,46 @@ def main() -> None:
         corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
         loaded.append((ds_name, corpus, queries, qrels))
 
-    output_dtype = np.float16 if args.qdrant_vector_dtype == "float16" else np.float32
+    # Resolve the dtype used for query embeddings:
+    # - If user sets float16/float32 explicitly, respect it.
+    # - If auto: try to detect from the existing Qdrant collection (prevents silent score drops),
+    #            otherwise fall back to float16 (preserves legacy default for new collections).
+    effective_qdrant_vector_dtype = str(args.qdrant_vector_dtype)
+    if effective_qdrant_vector_dtype == "auto":
+        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")
+        )
+        detected = None
+        if qdrant_url:
+            try:
+                from qdrant_client import QdrantClient
+
+                client = QdrantClient(
+                    url=qdrant_url,
+                    api_key=qdrant_api_key,
+                    prefer_grpc=bool(args.prefer_grpc),
+                    timeout=int(args.qdrant_timeout),
+                    check_compatibility=False,
+                )
+                detected = _detect_collection_vector_dtype(
+                    client=client, collection_name=str(args.collection)
+                )
+            except Exception:
+                detected = None
+        effective_qdrant_vector_dtype = detected or "float16"
+        if detected:
+            print(f"🔎 Detected Qdrant vector dtype for collection: {detected}")
+        else:
+            print("🔎 Could not detect Qdrant vector dtype; defaulting to float16")
+        sys.stdout.flush()
+
+    output_dtype = np.float16 if effective_qdrant_vector_dtype == "float16" else np.float32
     embedder = VisualEmbedder(
         model_name=args.model,
         batch_size=args.batch_size,
@@ -1105,9 +1414,6 @@ def main() -> None:
                     {"field": "original_height", "type": "integer"},
                     {"field": "resized_width", "type": "integer"},
                     {"field": "resized_height", "type": "integer"},
-                    {"field": "crop_empty_enabled", "type": "bool"},
-                    {"field": "crop_empty_remove_page_number", "type": "bool"},
-                    {"field": "crop_empty_percentage_to_remove", "type": "float"},
                     {"field": "num_tiles", "type": "integer"},
                     {"field": "tile_rows", "type": "integer"},
                     {"field": "tile_cols", "type": "integer"},
@@ -1118,6 +1424,28 @@ def main() -> None:
                     {"field": "source", "type": "keyword"},
                 ]
             )
+            # Keep schema minimal: only add crop-related indexes when cropping is enabled.
+            if bool(args.crop_empty):
+                payload_indexes.extend(
+                    [
+                        {"field": "crop_empty_enabled", "type": "bool"},
+                        {"field": "crop_empty_remove_page_number", "type": "bool"},
+                        {"field": "crop_empty_percentage_to_remove", "type": "float"},
+                        {"field": "cropped_url", "type": "keyword"},
+                        {"field": "cropped_width", "type": "integer"},
+                        {"field": "cropped_height", "type": "integer"},
+                    ]
+                )
+        # If we are recreating the collection, clear historical failure logs so they don't
+        # remove valid qrels during evaluation.
+        if bool(args.recreate):
+            for ds_name, _corpus, _queries, _qrels in selected:
+                p = _failed_log_path(collection_name=args.collection, dataset_name=ds_name)
+                try:
+                    if p.exists():
+                        p.unlink()
+                except Exception:
+                    pass
         for i, (ds_name, corpus, queries, _qrels) in enumerate(selected):
             print(f"Indexing {ds_name}: corpus_docs={len(corpus)} queries={len(queries)}")
             _index_beir_corpus(
@@ -1126,7 +1454,7 @@ def main() -> None:
                 embedder=embedder,
                 collection_name=args.collection,
                 prefer_grpc=args.prefer_grpc,
-                qdrant_vector_dtype=args.qdrant_vector_dtype,
+                qdrant_vector_dtype=effective_qdrant_vector_dtype,
                 recreate=bool(args.recreate and i == 0),
                 indexing_threshold=args.indexing_threshold,
                 batch_size=args.batch_size,
@@ -1148,6 +1476,7 @@ def main() -> None:
                 crop_empty_remove_page_number=bool(args.crop_empty_remove_page_number),
                 crop_empty_preserve_border_px=int(args.crop_empty_preserve_border_px),
                 crop_empty_uniform_std_threshold=float(args.crop_empty_uniform_std_threshold),
+                max_mean_pool_vectors=args.max_mean_pool_vectors,
                 no_cloudinary=bool(args.no_cloudinary),
                 cloudinary_folder=str(args.cloudinary_folder),
                 retry_failures=bool(args.retry_failures),
@@ -1160,9 +1489,15 @@ def main() -> None:
         dataset_index_failures: Dict[str, Dict[str, Any]] = {}
         dataset_counts: Dict[str, Dict[str, int]] = {}
         for ds_name, corpus, queries, _qrels in selected:
-            dataset_counts[ds_name] = {"corpus_docs": int(len(corpus)), "queries": int(len(queries)), "queries_eval": 0}
+            dataset_counts[ds_name] = {
+                "corpus_docs": int(len(corpus)),
+                "queries": int(len(queries)),
+                "queries_eval": 0,
+            }
             failed_path = _failed_log_path(collection_name=args.collection, dataset_name=ds_name)
-            failed_ids = _load_failed_union_ids(failed_path, dataset_name=ds_name, union_namespace=args.collection)
+            failed_ids = _load_failed_union_ids(
+                failed_path, dataset_name=ds_name, union_namespace=args.collection
+            )
             dataset_index_failures[ds_name] = {
                 "failed_log_path": str(failed_path),
                 "failed_ids_count": int(len(failed_ids)),
@@ -1178,7 +1513,7 @@ def main() -> None:
                 "collection": args.collection,
                 "model": args.model,
                 "torch_dtype": _torch_dtype_to_str(embedder.torch_dtype),
-                "qdrant_vector_dtype": args.qdrant_vector_dtype,
+                "qdrant_vector_dtype": effective_qdrant_vector_dtype,
                 "mode": args.mode,
                 "stage1_mode": args.stage1_mode if args.mode == "two_stage" else None,
                 "prefetch_k": args.prefetch_k if args.mode == "two_stage" else None,
@@ -1200,10 +1535,45 @@ def main() -> None:
         print(f"Wrote index-only report: {out_path}")
         return
 
+    if bool(args.ensure_in_ram):
+        try:
+            from visual_rag.qdrant_admin import QdrantAdmin
+
+            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")
+            )
+            admin = QdrantAdmin(
+                url=qdrant_url,
+                api_key=qdrant_api_key,
+                prefer_grpc=bool(args.prefer_grpc),
+                timeout=int(args.qdrant_timeout),
+            )
+            print(f"🧠 Ensuring collection in RAM (config): {args.collection}")
+            sys.stdout.flush()
+            _ = admin.ensure_collection_all_in_ram(
+                collection_name=str(args.collection),
+                timeout=int(args.qdrant_timeout),
+            )
+            print("✅ ensure-in-ram config applied")
+            sys.stdout.flush()
+        except Exception as e:
+            print(f"⚠️ ensure-in-ram failed: {type(e).__name__}: {e}")
+            sys.stdout.flush()
+
     retriever = MultiVectorRetriever(
         collection_name=args.collection,
         embedder=embedder,
-        qdrant_url=os.getenv("SIGIR_QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("QDRANT_URL"),
+        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")
@@ -1234,7 +1604,7 @@ def main() -> None:
             "collection": args.collection,
             "model": args.model,
             "torch_dtype": _torch_dtype_to_str(embedder.torch_dtype),
-            "qdrant_vector_dtype": args.qdrant_vector_dtype,
+            "qdrant_vector_dtype": effective_qdrant_vector_dtype,
             "mode": args.mode,
             "stage1_mode": args.stage1_mode if args.mode == "two_stage" else None,
             "prefetch_k": args.prefetch_k if args.mode == "two_stage" else None,
@@ -1271,13 +1641,25 @@ def main() -> None:
             f"(corpus_docs={len(corpus)}, queries={len(queries)}) "
             f"scope={args.evaluation_scope} "
             f"mode={args.mode}"
-            + (f", stage1_mode={args.stage1_mode}, prefetch_k={int(args.prefetch_k)}" if args.mode == "two_stage" else "")
-            + (f", stage1_k={int(args.stage1_k)}, stage2_k={int(args.stage2_k)}" if args.mode == "three_stage" else "")
+            + (
+                f", stage1_mode={args.stage1_mode}, prefetch_k={int(args.prefetch_k)}"
+                if args.mode == "two_stage"
+                else ""
+            )
+            + (
+                f", stage1_k={int(args.stage1_k)}, stage2_k={int(args.stage2_k)}"
+                if args.mode == "three_stage"
+                else ""
+            )
             + f", top_k={int(args.top_k)}"
         )
         sys.stdout.flush()
 
-        dataset_counts[ds_name] = {"corpus_docs": int(len(corpus)), "queries": int(len(queries)), "queries_eval": 0}
+        dataset_counts[ds_name] = {
+            "corpus_docs": int(len(corpus)),
+            "queries": int(len(queries)),
+            "queries_eval": 0,
+        }
         id_map: Dict[str, str] = {}
         for doc in corpus:
             source_doc_id = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
@@ -1298,11 +1680,21 @@ def main() -> None:
                 remapped_qrels[qid] = out_rels
 
         failed_path = _failed_log_path(collection_name=args.collection, dataset_name=ds_name)
-        failed_ids = _load_failed_union_ids(failed_path, dataset_name=ds_name, union_namespace=args.collection)
-        remapped_qrels, removed_rels = _remove_failed_from_qrels(remapped_qrels, failed_ids)
+        failed_ids_all = _load_failed_union_ids(
+            failed_path, dataset_name=ds_name, union_namespace=args.collection
+        )
+        # Only remove failed IDs that are actually missing in the current collection.
+        failed_ids_missing = _filter_failed_ids_to_missing(
+            qdrant_client=retriever.client,
+            collection_name=str(args.collection),
+            failed_ids=failed_ids_all,
+            timeout=int(args.qdrant_timeout),
+        )
+        remapped_qrels, removed_rels = _remove_failed_from_qrels(remapped_qrels, failed_ids_missing)
         dataset_index_failures[ds_name] = {
             "failed_log_path": str(failed_path),
-            "failed_ids_count": int(len(failed_ids)),
+            "failed_ids_count": int(len(failed_ids_all)),
+            "failed_ids_missing_count": int(len(failed_ids_missing)),
             "qrels_removed": int(removed_rels),
         }
 
@@ -1311,7 +1703,11 @@ def main() -> None:
             from qdrant_client.http import models as qmodels
 
             filter_obj = qmodels.Filter(
-                must=[qmodels.FieldCondition(key="dataset", match=qmodels.MatchValue(value=str(ds_name)))]
+                must=[
+                    qmodels.FieldCondition(
+                        key="dataset", match=qmodels.MatchValue(value=str(ds_name))
+                    )
+                ]
             )
 
         try:
@@ -1330,7 +1726,9 @@ def main() -> None:
                 drop_empty_queries=bool(args.drop_empty_queries),
                 filter_obj=filter_obj,
             )
-            dataset_counts[ds_name]["queries_eval"] = int(metrics_by_dataset[ds_name].get("num_queries_eval", 0))
+            dataset_counts[ds_name]["queries_eval"] = int(
+                metrics_by_dataset[ds_name].get("num_queries_eval", 0)
+            )
             ds_only_out = {
                 **_build_run_record(),
                 "dataset": str(ds_name),
@@ -1362,4 +1760,4 @@ def main() -> None:
 
 
 if __name__ == "__main__":
-    main()
+    main()

benchmarks/vidore_tatdqa_test/__init__.py

@@ -1,6 +1 @@
 __all__ = []
-
-
-
-
-

benchmarks/vidore_tatdqa_test/dataset_loader.py

@@ -1,8 +1,8 @@
 from __future__ import annotations
 
-from dataclasses import dataclass
 import hashlib
 import re
+from dataclasses import dataclass
 from typing import Any, Dict, Iterable, List, Mapping, Optional, Tuple
 
 
@@ -29,6 +29,7 @@ def _stable_uuid(text: str) -> str:
     hex_str = hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]
     return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:32]}"
 
+
 def paired_source_doc_id(row: Mapping[str, Any], idx: int) -> str:
     source_doc_id = _as_str(row.get("_id"))
     if source_doc_id:
@@ -55,7 +56,9 @@ def _normalize_qrels(qrels_rows: Iterable[Mapping[str, Any]]) -> Dict[str, Dict[
     qrels: Dict[str, Dict[str, int]] = {}
     for row in qrels_rows:
         qid = _as_str(row.get("query-id") or row.get("query_id") or row.get("qid"))
-        did = _as_str(row.get("corpus-id") or row.get("corpus_id") or row.get("doc_id") or row.get("did"))
+        did = _as_str(
+            row.get("corpus-id") or row.get("corpus_id") or row.get("doc_id") or row.get("did")
+        )
         score = row.get("score") or row.get("relevance") or row.get("label") or 0
         try:
             score_int = int(score)
@@ -63,6 +66,9 @@ def _normalize_qrels(qrels_rows: Iterable[Mapping[str, Any]]) -> Dict[str, Dict[
             score_int = 0
         if not qid or not did:
             continue
+        # Keep qrels compact and correct: score<=0 is non-relevant.
+        if score_int <= 0:
+            continue
         qrels.setdefault(qid, {})[_stable_uuid(did)] = score_int
     return qrels
 
@@ -70,7 +76,9 @@ def _normalize_qrels(qrels_rows: Iterable[Mapping[str, Any]]) -> Dict[str, Dict[
 def _expect_fields(obj: Any, required: List[str], context: str) -> None:
     missing = [k for k in required if k not in obj]
     if missing:
-        raise ValueError(f"{context}: missing required field(s): {missing}. Available: {list(obj.keys())}")
+        raise ValueError(
+            f"{context}: missing required field(s): {missing}. Available: {list(obj.keys())}"
+        )
 
 
 def _extract_beir_splits(ds: Any):
@@ -194,7 +202,9 @@ def _load_beir_from_separate_configs(dataset_name: str, config_names: List[str])
     return _first_split(corpus_ds), _first_split(queries_ds), _first_split(qrels_ds)
 
 
-def load_vidore_beir_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]]]:
+def load_vidore_beir_dataset(
+    dataset_name: str,
+) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]]]:
     try:
         from datasets import load_dataset
     except ImportError as e:
@@ -220,7 +230,6 @@ def load_vidore_beir_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Q
 
     last_err: Optional[Exception] = None
     extracted = None
-    used_name = None
     used_configs: List[str] = []
     for name_try in candidates:
         config_names = _get_config_names(name_try)
@@ -241,7 +250,6 @@ def load_vidore_beir_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Q
             if extracted is None:
                 extracted = _load_beir_from_separate_configs(name_try, config_names)
         if extracted is not None:
-            used_name = name_try
             break
 
     if extracted is None:
@@ -271,7 +279,9 @@ def load_vidore_beir_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Q
         doc_id = _stable_uuid(source_doc_id)
         image = row.get("image") or row.get("page_image") or row.get("document") or row.get("img")
         if image is None:
-            raise ValueError("corpus row: missing image field (tried image/page_image/document/img)")
+            raise ValueError(
+                "corpus row: missing image field (tried image/page_image/document/img)"
+            )
         payload = {
             **{
                 k: v
@@ -305,7 +315,9 @@ def load_vidore_beir_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Q
     return corpus_docs, queries, qrels
 
 
-def load_vidore_paired_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]]]:
+def load_vidore_paired_dataset(
+    dataset_name: str,
+) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]]]:
     """
     Load ViDoRe v1-style paired QA datasets.
 
@@ -347,7 +359,9 @@ def load_vidore_paired_dataset(dataset_name: str) -> Tuple[List[CorpusDoc], List
     return corpus_docs, queries, qrels
 
 
-def load_vidore_dataset_auto(dataset_name: str) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]], str]:
+def load_vidore_dataset_auto(
+    dataset_name: str,
+) -> Tuple[List[CorpusDoc], List[Query], Dict[str, Dict[str, int]], str]:
     """
     Auto-detect ViDoRe dataset format.
     Returns: (corpus, queries, qrels, protocol)
@@ -359,5 +373,3 @@ def load_vidore_dataset_auto(dataset_name: str) -> Tuple[List[CorpusDoc], List[Q
     except ValueError:
         corpus, queries, qrels = load_vidore_paired_dataset(dataset_name)
         return corpus, queries, qrels, "paired"
-
-

benchmarks/vidore_tatdqa_test/metrics.py

@@ -37,8 +37,3 @@ def recall_at_k(ranking: List[str], qrels: Dict[str, int], k: int) -> float:
         return 0.0
     retrieved = set(ranking[:k])
     return len(retrieved & relevant) / len(relevant)
-
-
-
-
-

benchmarks/vidore_tatdqa_test/run_qdrant.py

@@ -7,14 +7,17 @@ from typing import Any, Dict, List, Optional
 
 import numpy as np
 
+from benchmarks.vidore_tatdqa_test.dataset_loader import (
+    load_vidore_dataset_auto,
+    paired_doc_id,
+    paired_payload,
+)
+from benchmarks.vidore_tatdqa_test.metrics import mrr_at_k, ndcg_at_k, recall_at_k
 from visual_rag import VisualEmbedder
 from visual_rag.embedding.pooling import tile_level_mean_pooling
 from visual_rag.indexing.qdrant_indexer import QdrantIndexer
 from visual_rag.retrieval import MultiVectorRetriever
 
-from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_dataset_auto, paired_doc_id, paired_payload
-from benchmarks.vidore_tatdqa_test.metrics import ndcg_at_k, mrr_at_k, recall_at_k
-
 
 def _torch_dtype_to_str(dtype) -> str:
     if dtype is None:
@@ -144,7 +147,9 @@ def _index_corpus(
         pass
 
     def _upload(points: List[Dict[str, Any]]) -> int:
-        return indexer.upload_batch(points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event)
+        return indexer.upload_batch(
+            points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event
+        )
 
     executor = None
     futures = []
@@ -152,7 +157,8 @@ def _index_corpus(
 
     stop_event = threading.Event()
     if upload_workers and upload_workers > 0:
-        from concurrent.futures import ThreadPoolExecutor, wait as futures_wait, FIRST_EXCEPTION
+        from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor
+        from concurrent.futures import wait as futures_wait
 
         executor = ThreadPoolExecutor(max_workers=upload_workers)
 
@@ -226,9 +232,17 @@ def _index_corpus(
 
             for doc, emb, token_info in zip(batch, embeddings, token_infos):
                 if doc.image is None:
-                    raise ValueError("CorpusDoc.image is None. For paired datasets, use _index_paired_dataset().")
-                emb_np = emb.cpu().float().numpy() if hasattr(emb, "cpu") else np.array(emb, dtype=np.float32)
-                visual_indices = token_info.get("visual_token_indices") or list(range(emb_np.shape[0]))
+                    raise ValueError(
+                        "CorpusDoc.image is None. For paired datasets, use _index_paired_dataset()."
+                    )
+                emb_np = (
+                    emb.cpu().float().numpy()
+                    if hasattr(emb, "cpu")
+                    else np.array(emb, dtype=np.float32)
+                )
+                visual_indices = token_info.get("visual_token_indices") or list(
+                    range(emb_np.shape[0])
+                )
                 visual_embedding = emb_np[visual_indices].astype(np.float32)
 
                 n_rows = token_info.get("n_rows")
@@ -238,7 +252,9 @@ def _index_corpus(
                 else:
                     num_tiles = 13
 
-                tile_pooled = tile_level_mean_pooling(visual_embedding, num_tiles=num_tiles, patches_per_tile=64)
+                tile_pooled = tile_level_mean_pooling(
+                    visual_embedding, num_tiles=num_tiles, patches_per_tile=64
+                )
                 global_pooled = tile_pooled.mean(axis=0).astype(np.float32)
 
                 payload = {
@@ -270,8 +286,8 @@ def _index_corpus(
                     if pbar is not None:
                         pbar.set_postfix(
                             {
-                            "avg_s/doc": f"{avg_s_per_doc:.2f}",
-                            "last_s/doc": f"{last_s_per_doc:.2f}",
+                                "avg_s/doc": f"{avg_s_per_doc:.2f}",
+                                "last_s/doc": f"{last_s_per_doc:.2f}",
                                 "buffer": len(points_buffer),
                                 "enq": enqueued_docs,
                                 "upl": uploaded_docs,
@@ -343,7 +359,9 @@ def _index_paired_dataset(
         import torch
         from torch.utils.data import DataLoader
     except ImportError as e:
-        raise ImportError("torch is required. Install with: pip install visual-rag-toolkit[embedding]") from e
+        raise ImportError(
+            "torch is required. Install with: pip install visual-rag-toolkit[embedding]"
+        ) from e
 
     ds0 = load_dataset(dataset_name, split="test")
     cols = set(ds0.column_names)
@@ -389,7 +407,9 @@ def _index_paired_dataset(
         pass
 
     def _upload(points: List[Dict[str, Any]]) -> int:
-        return indexer.upload_batch(points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event)
+        return indexer.upload_batch(
+            points, delay_between_batches=0.0, wait=upsert_wait, stop_event=stop_event
+        )
 
     executor = None
     futures = []
@@ -397,7 +417,8 @@ def _index_paired_dataset(
 
     stop_event = threading.Event()
     if upload_workers and upload_workers > 0:
-        from concurrent.futures import ThreadPoolExecutor, wait as futures_wait, FIRST_EXCEPTION
+        from concurrent.futures import FIRST_EXCEPTION, ThreadPoolExecutor
+        from concurrent.futures import wait as futures_wait
 
         executor = ThreadPoolExecutor(max_workers=upload_workers)
 
@@ -446,7 +467,12 @@ def _index_paired_dataset(
                 dl_kwargs["pin_memory"] = bool(pin_memory and torch.cuda.is_available())
 
             data_loader = DataLoader(
-                _PairedHFDataset(dataset_name=dataset_name, split="test", total_docs=total_docs, image_col=image_col),
+                _PairedHFDataset(
+                    dataset_name=dataset_name,
+                    split="test",
+                    total_docs=total_docs,
+                    image_col=image_col,
+                ),
                 **dl_kwargs,
             )
             iterable = ((idxs, images, metas) for (idxs, images, metas) in data_loader)
@@ -457,7 +483,10 @@ def _index_paired_dataset(
                 for start in range(0, total_docs, batch_size):
                     batch = ds[start : start + batch_size]
                     images = batch[image_col]
-                    metas = [{k: batch[k][i] for k in batch.keys() if k != image_col} for i in range(len(images))]
+                    metas = [
+                        {k: batch[k][i] for k in batch.keys() if k != image_col}
+                        for i in range(len(images))
+                    ]
                     idxs = list(range(start, start + len(images)))
                     yield idxs, images, metas
 
@@ -501,15 +530,23 @@ def _index_paired_dataset(
                     **paired_payload(meta, int(idx)),
                 }
 
-                emb_np = emb.cpu().float().numpy() if hasattr(emb, "cpu") else np.array(emb, dtype=np.float32)
-                visual_indices = token_info.get("visual_token_indices") or list(range(emb_np.shape[0]))
+                emb_np = (
+                    emb.cpu().float().numpy()
+                    if hasattr(emb, "cpu")
+                    else np.array(emb, dtype=np.float32)
+                )
+                visual_indices = token_info.get("visual_token_indices") or list(
+                    range(emb_np.shape[0])
+                )
                 visual_embedding = emb_np[visual_indices].astype(np.float32)
 
                 n_rows = token_info.get("n_rows")
                 n_cols = token_info.get("n_cols")
                 num_tiles = int(n_rows) * int(n_cols) + 1 if n_rows and n_cols else 13
 
-                tile_pooled = tile_level_mean_pooling(visual_embedding, num_tiles=num_tiles, patches_per_tile=64)
+                tile_pooled = tile_level_mean_pooling(
+                    visual_embedding, num_tiles=num_tiles, patches_per_tile=64
+                )
                 global_pooled = tile_pooled.mean(axis=0).astype(np.float32)
 
                 points_buffer.append(
@@ -654,8 +691,14 @@ def main() -> None:
     grpc_group = parser.add_mutually_exclusive_group()
     grpc_group.add_argument("--prefer-grpc", dest="prefer_grpc", action="store_true", default=True)
     grpc_group.add_argument("--no-prefer-grpc", dest="prefer_grpc", action="store_false")
-    parser.add_argument("--index", action="store_true", help="Index corpus into Qdrant before evaluating")
-    parser.add_argument("--recreate", action="store_true", help="Delete and recreate the collection (implies --index)")
+    parser.add_argument(
+        "--index", action="store_true", help="Index corpus into Qdrant before evaluating"
+    )
+    parser.add_argument(
+        "--recreate",
+        action="store_true",
+        help="Delete and recreate the collection (implies --index)",
+    )
     parser.add_argument(
         "--indexing-threshold",
         type=int,
@@ -679,8 +722,19 @@ def main() -> None:
     parser.add_argument(
         "--stage1-mode",
         type=str,
-        default="tokens_vs_tiles",
-        choices=["pooled_query_vs_tiles", "tokens_vs_tiles", "pooled_query_vs_global"],
+        default="tokens_vs_standard_pooling",
+        choices=[
+            "pooled_query_vs_standard_pooling",
+            "tokens_vs_standard_pooling",
+            "pooled_query_vs_experimental_pooling",
+            "tokens_vs_experimental_pooling",
+            "pooled_query_vs_global",
+            # Backwards-compatible aliases
+            "pooled_query_vs_tiles",
+            "tokens_vs_tiles",
+            "pooled_query_vs_experimental",
+            "tokens_vs_experimental",
+        ],
     )
     parser.add_argument("--output", type=str, default="results/qdrant_vidore_tatdqa_test.json")
 
@@ -795,5 +849,3 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
-
-

benchmarks/vidore_tatdqa_test/sweep_eval.py

@@ -1,15 +1,15 @@
 import argparse
 import json
+import logging
 import os
 import time
-import logging
 from pathlib import Path
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional
 
 import numpy as np
 
 from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_dataset_auto
-from benchmarks.vidore_tatdqa_test.metrics import ndcg_at_k, mrr_at_k, recall_at_k
+from benchmarks.vidore_tatdqa_test.metrics import mrr_at_k, ndcg_at_k, recall_at_k
 from visual_rag import VisualEmbedder
 from visual_rag.retrieval import MultiVectorRetriever
 
@@ -24,6 +24,7 @@ def _maybe_load_dotenv() -> None:
     if Path(".env").exists():
         load_dotenv(".env")
 
+
 def _torch_dtype_to_str(dtype) -> str:
     if dtype is None:
         return "auto"
@@ -151,7 +152,9 @@ def _evaluate(
 
 
 def main() -> None:
-    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", force=True)
+    logging.basicConfig(
+        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s", force=True
+    )
 
     parser = argparse.ArgumentParser()
     parser.add_argument("--dataset", type=str, default="vidore/tatdqa_test")
@@ -165,12 +168,25 @@ def main() -> None:
         help="Torch dtype for model weights (default: auto; inferred from collection vector dtype when possible).",
     )
     parser.add_argument("--top-k", type=int, default=10)
-    parser.add_argument("--mode", type=str, default="two_stage", choices=["single_full", "two_stage"])
+    parser.add_argument(
+        "--mode", type=str, default="two_stage", choices=["single_full", "two_stage"]
+    )
     parser.add_argument(
         "--stage1-mode",
         type=str,
-        default="tokens_vs_tiles",
-        choices=["pooled_query_vs_tiles", "tokens_vs_tiles", "pooled_query_vs_global"],
+        default="tokens_vs_standard_pooling",
+        choices=[
+            "pooled_query_vs_standard_pooling",
+            "tokens_vs_standard_pooling",
+            "pooled_query_vs_experimental_pooling",
+            "tokens_vs_experimental_pooling",
+            "pooled_query_vs_global",
+            # Backwards-compatible aliases
+            "pooled_query_vs_tiles",
+            "tokens_vs_tiles",
+            "pooled_query_vs_experimental",
+            "tokens_vs_experimental",
+        ],
     )
     parser.add_argument(
         "--prefetch-ks",
@@ -278,7 +294,9 @@ def main() -> None:
     if args.query_batch_size and args.query_batch_size > 0:
         texts = [q.text for q in queries]
         logger.info(f"Pre-embedding {len(texts)} queries (batch={args.query_batch_size})...")
-        q_tensors = embedder.embed_queries(texts, batch_size=args.query_batch_size, show_progress=True)
+        q_tensors = embedder.embed_queries(
+            texts, batch_size=args.query_batch_size, show_progress=True
+        )
         precomputed_query_embeddings = [t.detach().cpu().float().numpy() for t in q_tensors]
         try:
             import torch
@@ -317,7 +335,11 @@ def main() -> None:
                     "max_queries": args.max_queries,
                     "sample_queries": args.sample_queries,
                     "sample_strategy": args.sample_strategy if args.sample_queries else None,
-                    "sample_seed": args.sample_seed if args.sample_queries and args.sample_strategy == "random" else None,
+                    "sample_seed": (
+                        args.sample_seed
+                        if args.sample_queries and args.sample_strategy == "random"
+                        else None
+                    ),
                     "metrics": metrics,
                 },
                 f,
@@ -340,7 +362,10 @@ def main() -> None:
             max_queries=args.max_queries,
             precomputed_query_embeddings=precomputed_query_embeddings,
         )
-        out_path = out_dir / f"{protocol}__two_stage__{args.stage1_mode}__prefetch{k}__top{args.top_k}.json"
+        out_path = (
+            out_dir
+            / f"{protocol}__two_stage__{args.stage1_mode}__prefetch{k}__top{args.top_k}.json"
+        )
         with open(out_path, "w") as f:
             json.dump(
                 {
@@ -356,7 +381,11 @@ def main() -> None:
                     "max_queries": args.max_queries,
                     "sample_queries": args.sample_queries,
                     "sample_strategy": args.sample_strategy if args.sample_queries else None,
-                    "sample_seed": args.sample_seed if args.sample_queries and args.sample_strategy == "random" else None,
+                    "sample_seed": (
+                        args.sample_seed
+                        if args.sample_queries and args.sample_strategy == "random"
+                        else None
+                    ),
                     "metrics": metrics,
                 },
                 f,
@@ -368,5 +397,3 @@ def main() -> None:
 
 if __name__ == "__main__":
     main()
-
-

demo/app.py

@@ -1,6 +1,5 @@
 """Main entry point for the Visual RAG Toolkit demo application."""
 
-import os
 import sys
 from pathlib import Path
 
@@ -11,7 +10,7 @@ _repo_root = _app_dir.parent
 if str(_repo_root) not in sys.path:
     sys.path.insert(0, str(_repo_root))
 
-from dotenv import load_dotenv
+from dotenv import load_dotenv  # noqa: E402
 
 # Load .env from the repo root (works both locally and in Docker)
 if (_repo_root / ".env").exists():
@@ -19,7 +18,7 @@ if (_repo_root / ".env").exists():
 if (_app_dir / ".env").exists():
     load_dotenv(_app_dir / ".env")
 
-import streamlit as st
+import streamlit as st  # noqa: E402
 
 st.set_page_config(
     page_title="Visual RAG Toolkit",
@@ -28,11 +27,11 @@ st.set_page_config(
     initial_sidebar_state="expanded",
 )
 
-from demo.ui.header import render_header
-from demo.ui.sidebar import render_sidebar
-from demo.ui.upload import render_upload_tab
-from demo.ui.playground import render_playground_tab
-from demo.ui.benchmark import render_benchmark_tab
+from demo.ui.benchmark import render_benchmark_tab  # noqa: E402
+from demo.ui.header import render_header  # noqa: E402
+from demo.ui.playground import render_playground_tab  # noqa: E402
+from demo.ui.sidebar import render_sidebar  # noqa: E402
+from demo.ui.upload import render_upload_tab  # noqa: E402
 
 
 def main():

demo/commands.py

@@ -43,13 +43,17 @@ def generate_python_index_code(config: Dict[str, Any]) -> str:
     prefer_grpc = config.get("prefer_grpc", True)
     crop_empty = config.get("crop_empty", False)
     max_docs = config.get("max_docs")
-    
+
     torch_dtype = config.get("torch_dtype", "float16")
     qdrant_dtype = config.get("qdrant_vector_dtype", "float16")
-    
-    torch_dtype_map = {"float16": "torch.float16", "float32": "torch.float32", "bfloat16": "torch.bfloat16"}
+
+    torch_dtype_map = {
+        "float16": "torch.float16",
+        "float32": "torch.float32",
+        "bfloat16": "torch.bfloat16",
+    }
     torch_dtype_val = torch_dtype_map.get(torch_dtype, "torch.float16")
-    
+
     code_lines = [
         "import os",
         "import torch",
@@ -61,96 +65,104 @@ def generate_python_index_code(config: Dict[str, Any]) -> str:
         f'COLLECTION = "{collection}"',
         f'MODEL = "{model}"',
         f"BATCH_SIZE = {batch_size}",
-        f'DATASETS = [{datasets_str}]',
-        f'TORCH_DTYPE = {torch_dtype_val}',
+        f"DATASETS = [{datasets_str}]",
+        f"TORCH_DTYPE = {torch_dtype_val}",
         f'QDRANT_DTYPE = "{qdrant_dtype}"',
     ]
-    
+
     if max_docs:
         code_lines.append(f"MAX_DOCS = {max_docs}  # Limit docs per dataset")
-    
-    code_lines.extend([
-        "",
-        "# Initialize embedder",
-        "embedder = VisualEmbedder(",
-        "    model_name=MODEL,",
-        "    torch_dtype=TORCH_DTYPE,",
-        ")",
-        "",
-        "# Initialize indexer",
-        "indexer = QdrantIndexer(",
-        '    url=os.getenv("QDRANT_URL"),',
-        '    api_key=os.getenv("QDRANT_API_KEY"),',
-        "    collection_name=COLLECTION,",
-        f"    prefer_grpc={prefer_grpc},",
-        "    vector_datatype=QDRANT_DTYPE,",
-        ")",
-        "",
-        "# Create collection",
-        f"indexer.create_collection(force_recreate={config.get('recreate', False)})",
-        'indexer.create_payload_indexes(fields=[',
-        '    {"field": "dataset", "type": "keyword"},',
-        '    {"field": "doc_id", "type": "keyword"},',
-        '    {"field": "source_doc_id", "type": "keyword"},',
-        "])",
-        "",
-        "# Index each dataset",
-        "for ds_name in DATASETS:",
-        "    print(f'Loading {ds_name}...')",
-        "    corpus, queries, qrels = load_vidore_beir_dataset(ds_name)",
-    ])
-    
+
+    code_lines.extend(
+        [
+            "",
+            "# Initialize embedder",
+            "embedder = VisualEmbedder(",
+            "    model_name=MODEL,",
+            "    torch_dtype=TORCH_DTYPE,",
+            ")",
+            "",
+            "# Initialize indexer",
+            "indexer = QdrantIndexer(",
+            '    url=os.getenv("QDRANT_URL"),',
+            '    api_key=os.getenv("QDRANT_API_KEY"),',
+            "    collection_name=COLLECTION,",
+            f"    prefer_grpc={prefer_grpc},",
+            "    vector_datatype=QDRANT_DTYPE,",
+            ")",
+            "",
+            "# Create collection",
+            f"indexer.create_collection(force_recreate={config.get('recreate', False)})",
+            "indexer.create_payload_indexes(fields=[",
+            '    {"field": "dataset", "type": "keyword"},',
+            '    {"field": "doc_id", "type": "keyword"},',
+            '    {"field": "source_doc_id", "type": "keyword"},',
+            "])",
+            "",
+            "# Index each dataset",
+            "for ds_name in DATASETS:",
+            "    print(f'Loading {ds_name}...')",
+            "    corpus, queries, qrels = load_vidore_beir_dataset(ds_name)",
+        ]
+    )
+
     if max_docs:
         code_lines.append("    corpus = corpus[:MAX_DOCS]  # Limit")
-    
-    code_lines.extend([
-        "    print(f'Indexing {len(corpus)} documents...')",
-        "",
-        "    for i in range(0, len(corpus), BATCH_SIZE):",
-        "        batch = corpus[i:i + BATCH_SIZE]",
-        "        images = [doc.image for doc in batch]",
-        "",
-        "        # Embed images",
-        "        embeddings, token_infos = embedder.embed_images(",
-        "            images, return_token_info=True",
-        "        )",
-        "",
-        "        # Build points with multi-vector representations",
-        "        points = []",
-        "        for doc, emb, info in zip(batch, embeddings, token_infos):",
-        "            emb_np = emb.cpu().numpy()",
-        "            visual_idx = info.get('visual_token_indices', range(len(emb_np)))",
-        "            visual_emb = emb_np[visual_idx]",
-        "",
-        "            tile_pooled = embedder.mean_pool_visual_embedding(visual_emb, info)",
-        "            experimental = embedder.experimental_pool_visual_embedding(",
-        "                visual_emb, info, mean_pool=tile_pooled",
-        "            )",
-        "            global_pooled = embedder.global_pool_from_mean_pool(tile_pooled)",
-        "",
-        "            points.append({",
-        '                "id": f"{ds_name}_{doc.doc_id}",',
-        '                "visual_embedding": visual_emb,',
-        '                "tile_pooled_embedding": tile_pooled,',
-        '                "experimental_pooled_embedding": experimental,',
-        '                "global_pooled_embedding": global_pooled,',
-        '                "metadata": {',
-        '                    "dataset": ds_name,',
-        '                    "doc_id": doc.doc_id,',
-        '                    "source_doc_id": doc.payload.get("source_doc_id"),',
-        "                },",
-        "            })",
-        "",
-        "        indexer.upload_batch(points)",
-        "        print(f'  Batch {i//BATCH_SIZE + 1}: {len(points)} uploaded')",
-        "",
-        '    print(f"Done: {ds_name}")',
-    ])
-    
+
+    code_lines.extend(
+        [
+            "    print(f'Indexing {len(corpus)} documents...')",
+            "",
+            "    for i in range(0, len(corpus), BATCH_SIZE):",
+            "        batch = corpus[i:i + BATCH_SIZE]",
+            "        images = [doc.image for doc in batch]",
+            "",
+            "        # Embed images",
+            "        embeddings, token_infos = embedder.embed_images(",
+            "            images, return_token_info=True",
+            "        )",
+            "",
+            "        # Build points with multi-vector representations",
+            "        points = []",
+            "        for doc, emb, info in zip(batch, embeddings, token_infos):",
+            "            emb_np = emb.cpu().numpy()",
+            "            visual_idx = info.get('visual_token_indices', range(len(emb_np)))",
+            "            visual_emb = emb_np[visual_idx]",
+            "",
+            "            tile_pooled = embedder.mean_pool_visual_embedding(visual_emb, info)",
+            "            experimental = embedder.experimental_pool_visual_embedding(",
+            "                visual_emb, info, mean_pool=tile_pooled",
+            "            )",
+            "            global_pooled = embedder.global_pool_from_mean_pool(tile_pooled)",
+            "",
+            "            points.append({",
+            '                "id": f"{ds_name}_{doc.doc_id}",',
+            '                "visual_embedding": visual_emb,',
+            '                "tile_pooled_embedding": tile_pooled,',
+            '                "experimental_pooled_embedding": experimental,',
+            '                "global_pooled_embedding": global_pooled,',
+            '                "metadata": {',
+            '                    "dataset": ds_name,',
+            '                    "doc_id": doc.doc_id,',
+            '                    "source_doc_id": doc.payload.get("source_doc_id"),',
+            "                },",
+            "            })",
+            "",
+            "        indexer.upload_batch(points)",
+            "        print(f'  Batch {i//BATCH_SIZE + 1}: {len(points)} uploaded')",
+            "",
+            '    print(f"Done: {ds_name}")',
+        ]
+    )
+
     if crop_empty:
-        code_lines.insert(3, "from visual_rag.preprocessing.crop_empty import crop_empty, CropEmptyConfig")
-        code_lines.insert(len(code_lines) - 20, "        # Note: Add crop_empty() preprocessing before embedding")
-    
+        code_lines.insert(
+            3, "from visual_rag.preprocessing.crop_empty import crop_empty, CropEmptyConfig"
+        )
+        code_lines.insert(
+            len(code_lines) - 20, "        # Note: Add crop_empty() preprocessing before embedding"
+        )
+
     return "\n".join(code_lines)
 
 
@@ -189,10 +201,14 @@ def generate_python_eval_code(config: Dict[str, Any]) -> str:
     scope = config.get("evaluation_scope", "union")
     prefer_grpc = config.get("prefer_grpc", True)
     torch_dtype = config.get("torch_dtype", "float16")
-    
-    torch_dtype_map = {"float16": "torch.float16", "float32": "torch.float32", "bfloat16": "torch.bfloat16"}
+
+    torch_dtype_map = {
+        "float16": "torch.float16",
+        "float32": "torch.float32",
+        "bfloat16": "torch.bfloat16",
+    }
     torch_dtype_val = torch_dtype_map.get(torch_dtype, "torch.float16")
-    
+
     code_lines = [
         "import os",
         "import torch",
@@ -204,7 +220,7 @@ def generate_python_eval_code(config: Dict[str, Any]) -> str:
         f'COLLECTION = "{collection}"',
         f'MODEL = "{model}"',
         f"TOP_K = {top_k}",
-        f'DATASETS = [{datasets_str}]',
+        f"DATASETS = [{datasets_str}]",
         f"TORCH_DTYPE = {torch_dtype_val}",
         "",
         "# Initialize clients",
@@ -227,108 +243,122 @@ def generate_python_eval_code(config: Dict[str, Any]) -> str:
         ")",
         "",
     ]
-    
+
     if mode == "single_full":
-        code_lines.extend([
-            "# Single-stage full retrieval",
-            "def search(query: str):",
-            "    query_embedding = embedder.embed_query(query)",
-            "    return retriever.search_single_stage(",
-            "        query_embedding=query_embedding,",
-            f"        limit={top_k},",
-            '        vector_name="initial",',
-            "    )",
-        ])
+        code_lines.extend(
+            [
+                "# Single-stage full retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="initial",',
+                "    )",
+            ]
+        )
     elif mode == "single_tiles":
-        code_lines.extend([
-            "# Single-stage tiles retrieval",
-            "def search(query: str):",
-            "    query_embedding = embedder.embed_query(query)",
-            "    return retriever.search_single_stage(",
-            "        query_embedding=query_embedding,",
-            f"        limit={top_k},",
-            '        vector_name="mean_pooling",',
-            "    )",
-        ])
+        code_lines.extend(
+            [
+                "# Single-stage tiles retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="mean_pooling",',
+                "    )",
+            ]
+        )
     elif mode == "single_global":
-        code_lines.extend([
-            "# Single-stage global retrieval",
-            "def search(query: str):",
-            "    query_embedding = embedder.embed_query(query)",
-            "    return retriever.search_single_stage(",
-            "        query_embedding=query_embedding,",
-            f"        limit={top_k},",
-            '        vector_name="global_pooling",',
-            "    )",
-        ])
+        code_lines.extend(
+            [
+                "# Single-stage global retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="global_pooling",',
+                "    )",
+            ]
+        )
     elif mode == "two_stage":
         prefetch_k = config.get("prefetch_k", 256)
-        stage1_mode = config.get("stage1_mode", "tokens_vs_tiles")
-        code_lines.extend([
-            "# Two-stage retrieval",
-            "from visual_rag.retrieval import TwoStageRetriever",
-            "",
-            "two_stage = TwoStageRetriever(",
-            "    client=client,",
-            "    collection_name=COLLECTION,",
-            "    embedder=embedder,",
-            ")",
-            "",
-            "def search(query: str):",
-            "    query_embedding = embedder.embed_query(query)",
-            "    return two_stage.search(",
-            "        query_embedding=query_embedding,",
-            f"        prefetch_limit={prefetch_k},",
-            f"        limit={top_k},",
-            f'        stage1_mode="{stage1_mode}",',
-            "    )",
-        ])
+        stage1_mode = config.get("stage1_mode", "tokens_vs_standard_pooling")
+        code_lines.extend(
+            [
+                "# Two-stage retrieval",
+                "from visual_rag.retrieval import TwoStageRetriever",
+                "",
+                "two_stage = TwoStageRetriever(",
+                "    client=client,",
+                "    collection_name=COLLECTION,",
+                "    embedder=embedder,",
+                ")",
+                "",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return two_stage.search(",
+                "        query_embedding=query_embedding,",
+                f"        prefetch_limit={prefetch_k},",
+                f"        limit={top_k},",
+                f'        stage1_mode="{stage1_mode}",',
+                "    )",
+            ]
+        )
     elif mode == "three_stage":
         stage1_k = config.get("stage1_k", 1000)
         stage2_k = config.get("stage2_k", 300)
-        code_lines.extend([
-            "# Three-stage retrieval",
-            "from visual_rag.retrieval import ThreeStageRetriever",
-            "",
-            "three_stage = ThreeStageRetriever(",
-            "    client=client,",
-            "    collection_name=COLLECTION,",
-            "    embedder=embedder,",
-            ")",
-            "",
-            "def search(query: str):",
-            "    query_embedding = embedder.embed_query(query)",
-            "    return three_stage.search(",
-            "        query_embedding=query_embedding,",
-            f"        stage1_limit={stage1_k},",
-            f"        stage2_limit={stage2_k},",
-            f"        limit={top_k},",
-            "    )",
-        ])
-    
+        code_lines.extend(
+            [
+                "# Three-stage retrieval",
+                "from visual_rag.retrieval import ThreeStageRetriever",
+                "",
+                "three_stage = ThreeStageRetriever(",
+                "    client=client,",
+                "    collection_name=COLLECTION,",
+                "    embedder=embedder,",
+                ")",
+                "",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return three_stage.search(",
+                "        query_embedding=query_embedding,",
+                f"        stage1_limit={stage1_k},",
+                f"        stage2_limit={stage2_k},",
+                f"        limit={top_k},",
+                "    )",
+            ]
+        )
+
     if scope == "per_dataset":
-        code_lines.extend([
-            "",
-            "# Per-dataset filtering",
-            "from qdrant_client.models import Filter, FieldCondition, MatchValue",
+        code_lines.extend(
+            [
+                "",
+                "# Per-dataset filtering",
+                "from qdrant_client.models import Filter, FieldCondition, MatchValue",
+                "",
+                'def search_dataset(query: str, dataset: str = "vidore/esg_reports_v2"):',
+                "    query_embedding = embedder.embed_query(query)",
+                "    dataset_filter = Filter(",
+                "        must=[FieldCondition(",
+                '            key="dataset",',
+                "            match=MatchValue(value=dataset),",
+                "        )]",
+                "    )",
+                "    # Add filter to your search call",
+            ]
+        )
+
+    code_lines.extend(
+        [
             "",
-            'def search_dataset(query: str, dataset: str = "vidore/esg_reports_v2"):',
-            "    query_embedding = embedder.embed_query(query)",
-            "    dataset_filter = Filter(",
-            "        must=[FieldCondition(",
-            '            key="dataset",',
-            "            match=MatchValue(value=dataset),",
-            "        )]",
-            "    )",
-            "    # Add filter to your search call",
-        ])
-    
-    code_lines.extend([
-        "",
-        '# Example usage',
-        'results = search("What is the company revenue?")',
-        'for r in results:',
-        '    print(f"Score: {r.score:.4f}, Doc: {r.payload.get(\'doc_id\')}")',
-    ])
-    
+            "# Example usage",
+            'results = search("What is the company revenue?")',
+            "for r in results:",
+            "    print(f\"Score: {r.score:.4f}, Doc: {r.payload.get('doc_id')}\")",
+        ]
+    )
+
     return "\n".join(code_lines)

demo/config.py

@@ -3,6 +3,7 @@
 AVAILABLE_MODELS = [
     "vidore/colpali-v1.3",
     "vidore/colSmol-500M",
+    "vidore/colqwen2.5-v0.2",
 ]
 
 BENCHMARK_DATASETS = [
@@ -26,9 +27,9 @@ RETRIEVAL_MODES = [
 ]
 
 STAGE1_MODES = [
-    "tokens_vs_tiles",
-    "tokens_vs_experimental",
-    "pooled_query_vs_tiles",
-    "pooled_query_vs_experimental",
+    "tokens_vs_standard_pooling",
+    "tokens_vs_experimental_pooling",
+    "pooled_query_vs_standard_pooling",
+    "pooled_query_vs_experimental_pooling",
     "pooled_query_vs_global",
 ]

demo/download_models.py

@@ -6,7 +6,6 @@ avoiding download delays during container startup.
 """
 
 import os
-import sys
 
 os.environ.setdefault("HF_HOME", "/app/.cache/huggingface")
 os.environ.setdefault("TRANSFORMERS_CACHE", "/app/.cache/huggingface")
@@ -14,20 +13,22 @@ os.environ.setdefault("TRANSFORMERS_CACHE", "/app/.cache/huggingface")
 MODELS_TO_DOWNLOAD = [
     "vidore/colpali-v1.3",
     "vidore/colSmol-500M",
+    "vidore/colqwen2.5-v0.2",
 ]
 
+
 def download_colpali_models():
     """Download ColPali models and their processors."""
     print("=" * 60)
     print("Downloading ColPali models for Visual RAG Toolkit")
     print("=" * 60)
-    
+
     try:
         from colpali_engine.models import ColPali, ColPaliProcessor
     except ImportError:
         print("[WARN] colpali-engine not installed, trying transformers directly")
         from transformers import AutoModel, AutoProcessor
-        
+
         for model_name in MODELS_TO_DOWNLOAD:
             print(f"\n[INFO] Downloading model: {model_name}")
             try:
@@ -37,12 +38,19 @@ def download_colpali_models():
             except Exception as e:
                 print(f"[WARN] Could not download {model_name}: {e}")
         return
-    
+
     for model_name in MODELS_TO_DOWNLOAD:
         print(f"\n[INFO] Downloading model: {model_name}")
         try:
-            if "colsmol" in model_name.lower():
+            model_lower = model_name.lower()
+            if "colqwen2.5" in model_lower or "colqwen2_5" in model_lower:
+                from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor
+
+                ColQwen2_5.from_pretrained(model_name, trust_remote_code=True)
+                ColQwen2_5_Processor.from_pretrained(model_name, trust_remote_code=True)
+            elif "colqwen" in model_lower:
                 from colpali_engine.models import ColQwen2, ColQwen2Processor
+
                 ColQwen2.from_pretrained(model_name, trust_remote_code=True)
                 ColQwen2Processor.from_pretrained(model_name, trust_remote_code=True)
             else:
@@ -53,6 +61,7 @@ def download_colpali_models():
             print(f"[WARN] Could not download {model_name} with colpali-engine: {e}")
             try:
                 from transformers import AutoModel, AutoProcessor
+
                 AutoModel.from_pretrained(model_name, trust_remote_code=True)
                 AutoProcessor.from_pretrained(model_name, trust_remote_code=True)
                 print(f"[OK] Downloaded via transformers: {model_name}")
@@ -63,9 +72,9 @@ def download_colpali_models():
 def main():
     print(f"[INFO] HF_HOME: {os.environ.get('HF_HOME', 'not set')}")
     print(f"[INFO] Cache dir: {os.environ.get('TRANSFORMERS_CACHE', 'not set')}")
-    
+
     download_colpali_models()
-    
+
     print("\n" + "=" * 60)
     print("Model download complete!")
     print("=" * 60)

demo/evaluation.py

@@ -13,12 +13,11 @@ import streamlit as st
 import torch
 from qdrant_client.models import FieldCondition, Filter, MatchValue
 
-from visual_rag import VisualEmbedder
-from visual_rag.retrieval import MultiVectorRetriever
 from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
-from benchmarks.vidore_tatdqa_test.metrics import ndcg_at_k, mrr_at_k, recall_at_k
+from benchmarks.vidore_tatdqa_test.metrics import mrr_at_k, ndcg_at_k, recall_at_k
 from demo.qdrant_utils import get_qdrant_credentials
-
+from visual_rag import VisualEmbedder
+from visual_rag.retrieval import MultiVectorRetriever
 
 TORCH_DTYPE_MAP = {
     "float16": torch.float16,
@@ -36,7 +35,9 @@ def _stable_uuid(text: str) -> str:
     return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:32]}"
 
 
-def _union_point_id(*, dataset_name: str, source_doc_id: str, union_namespace: Optional[str]) -> str:
+def _union_point_id(
+    *, dataset_name: str, source_doc_id: str, union_namespace: Optional[str]
+) -> str:
     """Generate union point ID (same as benchmark script)."""
     ns = f"{union_namespace}::{dataset_name}" if union_namespace else dataset_name
     return _stable_uuid(f"{ns}::{source_doc_id}")
@@ -57,7 +58,7 @@ def _remap_qrels_to_union_ids(
             source_doc_id=source_doc_id,
             union_namespace=collection_name,
         )
-    
+
     remapped: Dict[str, Dict[str, int]] = {}
     for qid, rels in qrels.items():
         out_rels: Dict[str, int] = {}
@@ -72,7 +73,7 @@ def _remap_qrels_to_union_ids(
 
 def get_doc_id_from_result(r: Dict[str, Any], use_original: bool = True) -> str:
     """Extract document ID from search result.
-    
+
     Args:
         r: Search result dict with 'id' and 'payload'
         use_original: If True, prefer original doc_id for matching with qrels.
@@ -88,58 +89,54 @@ def get_doc_id_from_result(r: Dict[str, Any], use_original: bool = True) -> str:
             or str(r.get("id", ""))
         )
     else:
-        doc_id = (
-            payload.get("union_doc_id")
-            or str(r.get("id", ""))
-            or payload.get("doc_id")
-        )
+        doc_id = payload.get("union_doc_id") or str(r.get("id", "")) or payload.get("doc_id")
     return str(doc_id)
 
 
 def run_evaluation_with_ui(config: Dict[str, Any]):
     st.divider()
-    
+
     print("=" * 60)
     print("[EVAL] Starting evaluation via UI")
     print("=" * 60)
-    
+
     url, api_key = get_qdrant_credentials()
     if not url:
         st.error("QDRANT_URL not configured")
         return
-    
+
     datasets = config.get("datasets", [])
     collection = config["collection"]
     model = config.get("model", "vidore/colpali-v1.3")
     mode = config.get("mode", "single_full")
     top_k = config.get("top_k", 100)
     prefetch_k = config.get("prefetch_k", 256)
-    stage1_mode = config.get("stage1_mode", "tokens_vs_tiles")
+    stage1_mode = config.get("stage1_mode", "tokens_vs_standard_pooling")
     stage1_k = config.get("stage1_k", 1000)
     stage2_k = config.get("stage2_k", 300)
     prefer_grpc = config.get("prefer_grpc", True)
     torch_dtype = config.get("torch_dtype", "float16")
     evaluation_scope = config.get("evaluation_scope", "union")
-    
-    print(f"[EVAL] ═══════════════════════════════════════════════════")
+
+    print("[EVAL] ═══════════════════════════════════════════════════")
     print(f"[EVAL] Collection: {collection}")
     print(f"[EVAL] Model: {model}")
     print(f"[EVAL] Mode: {mode}, Scope: {evaluation_scope}")
     print(f"[EVAL] Datasets: {datasets}")
     print(f"[EVAL] Query embedding dtype: {torch_dtype} (vectors already indexed)")
-    print(f"[EVAL] ═══════════════════════════════════════════════════")
-    
+    print("[EVAL] ═══════════════════════════════════════════════════")
+
     phase1_container = st.container()
     phase2_container = st.container()
     phase3_container = st.container()
     results_container = st.container()
-    
+
     try:
         with phase1_container:
             st.markdown("##### 🤖 Phase 1: Loading Model")
             model_status = st.empty()
             model_status.info(f"Loading `{model.split('/')[-1]}`...")
-            
+
             print(f"[EVAL] Loading embedder: {model}")
             torch_dtype_obj = TORCH_DTYPE_MAP.get(torch_dtype, torch.float16)
             qdrant_dtype = config.get("qdrant_vector_dtype", "float16")
@@ -150,13 +147,15 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 output_dtype=output_dtype_obj,
             )
             embedder._load_model()
-            print(f"[EVAL] Embedder loaded (torch_dtype={torch_dtype}, output_dtype={qdrant_dtype})")
-            
+            print(
+                f"[EVAL] Embedder loaded (torch_dtype={torch_dtype}, output_dtype={qdrant_dtype})"
+            )
+
             model_status.success(f"✅ Model `{model.split('/')[-1]}` loaded")
-            
+
             retriever_status = st.empty()
             retriever_status.info(f"Connecting to collection `{collection}`...")
-            
+
             print(f"[EVAL] Connecting to Qdrant collection: {collection}")
             retriever = MultiVectorRetriever(
                 collection_name=collection,
@@ -166,31 +165,33 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 prefer_grpc=prefer_grpc,
                 embedder=embedder,
             )
-            print(f"[EVAL] Connected to Qdrant")
+            print("[EVAL] Connected to Qdrant")
             retriever_status.success(f"✅ Connected to `{collection}`")
-        
+
         with phase2_container:
             st.markdown("##### 📚 Phase 2: Loading Datasets")
-            
+
             dataset_data = {}
             total_queries = 0
             max_queries_per_ds = config.get("max_queries")
-            
+
             for ds_name in datasets:
                 ds_status = st.empty()
                 ds_short = ds_name.split("/")[-1]
                 ds_status.info(f"Loading `{ds_short}`...")
-                
+
                 print(f"[EVAL] Loading dataset: {ds_name}")
                 corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
-                
+
                 print(f"[EVAL] Remapping qrels to union_doc_id format for collection={collection}")
                 remapped_qrels = _remap_qrels_to_union_ids(qrels, corpus, ds_name, collection)
-                print(f"[EVAL] Remapped {len(qrels)} -> {len(remapped_qrels)} queries with valid rels")
-                
+                print(
+                    f"[EVAL] Remapped {len(qrels)} -> {len(remapped_qrels)} queries with valid rels"
+                )
+
                 if evaluation_scope == "per_dataset" and max_queries_per_ds:
                     queries = queries[:max_queries_per_ds]
-                
+
                 dataset_data[ds_name] = {
                     "queries": queries,
                     "qrels": remapped_qrels,
@@ -199,69 +200,85 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 total_queries += len(queries)
                 print(f"[EVAL] Loaded {ds_name}: {len(corpus)} docs, {len(queries)} queries")
                 ds_status.success(f"✅ `{ds_short}`: {len(corpus)} docs, {len(queries)} queries")
-            
-            if evaluation_scope == "union" and max_queries_per_ds and max_queries_per_ds < total_queries:
+
+            if (
+                evaluation_scope == "union"
+                and max_queries_per_ds
+                and max_queries_per_ds < total_queries
+            ):
                 total_queries = max_queries_per_ds
                 print(f"[EVAL] Will limit to {total_queries} total queries (union mode)")
-            
+
             embed_status = st.empty()
-            embed_status.info(f"Embedding queries...")
-        
+            embed_status.info("Embedding queries...")
+
         with phase3_container:
             st.markdown("##### 🎯 Phase 3: Running Evaluation")
-            
+
             metrics_collectors = {
-                "ndcg@5": [], "ndcg@10": [],
-                "recall@5": [], "recall@10": [],
-                "mrr@5": [], "mrr@10": [],
+                "ndcg@5": [],
+                "ndcg@10": [],
+                "recall@5": [],
+                "recall@10": [],
+                "mrr@5": [],
+                "mrr@10": [],
             }
             latencies = []
             log_lines = []
             metrics_by_dataset = {}
-            
+
             if evaluation_scope == "per_dataset":
                 overall_progress = st.progress(0.0)
                 datasets_done = 0
-                
+
                 for ds_name, ds_info in dataset_data.items():
                     ds_short = ds_name.split("/")[-1]
                     st.markdown(f"**Evaluating `{ds_short}`**")
-                    
+
                     queries = ds_info["queries"]
                     qrels = ds_info["qrels"]
-                    
+
                     if not queries:
                         continue
-                    
+
                     print(f"[EVAL] Embedding {len(queries)} queries for {ds_short}...")
                     query_texts = [q.text for q in queries]
                     query_embeddings = embedder.embed_queries(query_texts, show_progress=False)
                     print(f"[EVAL] Queries embedded for {ds_short}")
-                    
+
                     ds_filter = Filter(
                         must=[FieldCondition(key="dataset", match=MatchValue(value=ds_name))]
                     )
                     print(f"[EVAL] Using filter: dataset={ds_name}")
-                    
+
                     progress_bar = st.progress(0.0)
                     eval_status = st.empty()
                     log_area = st.empty()
-                    
-                    ds_metrics = {"ndcg@5": [], "ndcg@10": [], "recall@5": [], "recall@10": [], "mrr@5": [], "mrr@10": []}
+
+                    ds_metrics = {
+                        "ndcg@5": [],
+                        "ndcg@10": [],
+                        "recall@5": [],
+                        "recall@10": [],
+                        "mrr@5": [],
+                        "mrr@10": [],
+                    }
                     ds_latencies = []
                     ds_log_lines = []
-                    
+
                     eval_status.info(f"Evaluating {len(queries)} queries...")
-                    print(f"[EVAL] Starting per-dataset evaluation: {ds_short}, {len(queries)} queries")
-                    
+                    print(
+                        f"[EVAL] Starting per-dataset evaluation: {ds_short}, {len(queries)} queries"
+                    )
+
                     for i, (q, qemb) in enumerate(zip(queries, query_embeddings)):
                         start = time.time()
-                        
+
                         if isinstance(qemb, torch.Tensor):
                             qemb_np = qemb.detach().cpu().numpy()
                         else:
-                            qemb_np = qemb.numpy() if hasattr(qemb, 'numpy') else np.array(qemb)
-                        
+                            qemb_np = qemb.numpy() if hasattr(qemb, "numpy") else np.array(qemb)
+
                         results = retriever.search_embedded(
                             query_embedding=qemb_np,
                             top_k=max(100, top_k),
@@ -274,61 +291,79 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                         )
                         ds_latencies.append((time.time() - start) * 1000)
                         latencies.append(ds_latencies[-1])
-                        
+
                         ranking = [str(r["id"]) for r in results]
                         rels = qrels.get(q.query_id, {})
-                        
+
                         if i == 0:
                             print(f"[EVAL] First query for {ds_short} - query_id: {q.query_id}")
                             print(f"[EVAL] Top 3 retrieved doc_ids: {ranking[:3]}")
                             print(f"[EVAL] Expected doc_ids (qrels): {list(rels.keys())[:3]}")
-                            print(f"[EVAL] qrels has {len(qrels)} queries, this query in qrels: {q.query_id in qrels}")
+                            print(
+                                f"[EVAL] qrels has {len(qrels)} queries, this query in qrels: {q.query_id in qrels}"
+                            )
                             if results:
                                 r0 = results[0]
-                                print(f"[EVAL] Sample result payload keys: {list(r0.get('payload', {}).keys())}")
+                                print(
+                                    f"[EVAL] Sample result payload keys: {list(r0.get('payload', {}).keys())}"
+                                )
                                 p = r0.get("payload", {})
-                                print(f"[EVAL] Sample payload doc_id={p.get('doc_id')}, union_doc_id={p.get('union_doc_id')}, source_doc_id={p.get('source_doc_id')}")
+                                print(
+                                    f"[EVAL] Sample payload doc_id={p.get('doc_id')}, union_doc_id={p.get('union_doc_id')}, source_doc_id={p.get('source_doc_id')}"
+                                )
                             has_match = any(rid in rels for rid in ranking[:10])
                             print(f"[EVAL] Any match in top-10? {has_match}")
-                        
+
                         for k_name, k_val in [("ndcg@5", 5), ("ndcg@10", 10)]:
                             ds_metrics[k_name].append(ndcg_at_k(ranking, rels, k=k_val))
                         for k_name, k_val in [("recall@5", 5), ("recall@10", 10)]:
                             ds_metrics[k_name].append(recall_at_k(ranking, rels, k=k_val))
                         for k_name, k_val in [("mrr@5", 5), ("mrr@10", 10)]:
                             ds_metrics[k_name].append(mrr_at_k(ranking, rels, k=k_val))
-                        
+
                         progress = (i + 1) / len(queries)
                         progress_bar.progress(progress)
-                        eval_status.info(f"🎯 {i+1}/{len(queries)} ({int(progress*100)}%) — latency: {np.mean(ds_latencies):.0f}ms")
-                        
+                        eval_status.info(
+                            f"🎯 {i+1}/{len(queries)} ({int(progress*100)}%) — latency: {np.mean(ds_latencies):.0f}ms"
+                        )
+
                         log_interval = max(5, len(queries) // 10)
                         if (i + 1) % log_interval == 0 and i > 0:
                             cur_ndcg = np.mean(ds_metrics["ndcg@10"])
-                            cur_lat = np.mean(ds_latencies[1:]) if len(ds_latencies) > 1 else ds_latencies[0]
-                            ds_log_lines.append(f"[{i+1}/{len(queries)}] NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms")
+                            cur_lat = (
+                                np.mean(ds_latencies[1:])
+                                if len(ds_latencies) > 1
+                                else ds_latencies[0]
+                            )
+                            ds_log_lines.append(
+                                f"[{i+1}/{len(queries)}] NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms"
+                            )
                             log_area.code("\n".join(ds_log_lines[-5:]), language="text")
-                            print(f"[EVAL] {ds_short} {i+1}/{len(queries)}: NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms")
-                    
+                            print(
+                                f"[EVAL] {ds_short} {i+1}/{len(queries)}: NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms"
+                            )
+
                     progress_bar.progress(1.0)
                     ds_final = {k: float(np.mean(v)) for k, v in ds_metrics.items()}
                     ds_final["avg_latency_ms"] = float(np.mean(ds_latencies))
                     ds_final["num_queries"] = len(queries)
                     metrics_by_dataset[ds_name] = ds_final
-                    
+
                     for k, v in ds_metrics.items():
                         metrics_collectors[k].extend(v)
-                    
-                    eval_status.success(f"✅ `{ds_short}`: NDCG@10={ds_final['ndcg@10']:.4f}, latency={ds_final['avg_latency_ms']:.0f}ms")
+
+                    eval_status.success(
+                        f"✅ `{ds_short}`: NDCG@10={ds_final['ndcg@10']:.4f}, latency={ds_final['avg_latency_ms']:.0f}ms"
+                    )
                     print(f"[EVAL] {ds_short} DONE: NDCG@10={ds_final['ndcg@10']:.4f}")
-                    
+
                     datasets_done += 1
                     overall_progress.progress(datasets_done / len(datasets))
-                
+
                 overall_progress.progress(1.0)
-                embed_status.success(f"✅ All queries embedded")
+                embed_status.success("✅ All queries embedded")
                 total_queries = sum(d["num_queries"] for d in metrics_by_dataset.values())
-            
+
             else:
                 all_queries = []
                 all_qrels = {}
@@ -336,7 +371,7 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                     all_queries.extend(ds_info["queries"])
                     for qid, rels in ds_info["qrels"].items():
                         all_qrels[qid] = rels
-                
+
                 sample_qrel_keys = list(all_qrels.keys())[:3]
                 sample_doc_ids = []
                 for qid in sample_qrel_keys:
@@ -344,33 +379,33 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 print(f"[EVAL] all_qrels built: {len(all_qrels)} queries")
                 print(f"[EVAL] Sample qrel query_ids: {sample_qrel_keys}")
                 print(f"[EVAL] Sample qrel doc_ids: {sample_doc_ids[:5]}")
-                
+
                 max_q = config.get("max_queries")
                 if max_q and max_q < len(all_queries):
                     all_queries = all_queries[:max_q]
                 total_queries = len(all_queries)
-                
+
                 print(f"[EVAL] Embedding {total_queries} queries (union mode)...")
                 query_texts = [q.text for q in all_queries]
                 query_embeddings = embedder.embed_queries(query_texts, show_progress=False)
-                print(f"[EVAL] Queries embedded")
+                print("[EVAL] Queries embedded")
                 embed_status.success(f"✅ {total_queries} queries embedded")
-                
+
                 progress_bar = st.progress(0.0)
                 eval_status = st.empty()
                 log_area = st.empty()
-                
+
                 eval_status.info(f"Evaluating {total_queries} queries in `{mode}` mode...")
                 print(f"[EVAL] Starting union evaluation: {total_queries} queries, mode={mode}")
-                
+
                 for i, (q, qemb) in enumerate(zip(all_queries, query_embeddings)):
                     start = time.time()
-                    
+
                     if isinstance(qemb, torch.Tensor):
                         qemb_np = qemb.detach().cpu().numpy()
                     else:
-                        qemb_np = qemb.numpy() if hasattr(qemb, 'numpy') else np.array(qemb)
-                    
+                        qemb_np = qemb.numpy() if hasattr(qemb, "numpy") else np.array(qemb)
+
                     results = retriever.search_embedded(
                         query_embedding=qemb_np,
                         top_k=max(100, top_k),
@@ -381,52 +416,64 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                         stage2_k=stage2_k,
                     )
                     latencies.append((time.time() - start) * 1000)
-                    
+
                     ranking = [str(r["id"]) for r in results]
                     rels = all_qrels.get(q.query_id, {})
-                    
+
                     if i == 0:
                         print(f"[EVAL] First query - query_id: {q.query_id}")
                         print(f"[EVAL] Top 3 retrieved doc_ids: {ranking[:3]}")
                         print(f"[EVAL] Expected doc_ids (qrels): {list(rels.keys())[:3]}")
-                        print(f"[EVAL] all_qrels has {len(all_qrels)} queries, this query in qrels: {q.query_id in all_qrels}")
+                        print(
+                            f"[EVAL] all_qrels has {len(all_qrels)} queries, this query in qrels: {q.query_id in all_qrels}"
+                        )
                         if results:
                             r0 = results[0]
-                            print(f"[EVAL] Sample result payload keys: {list(r0.get('payload', {}).keys())}")
+                            print(
+                                f"[EVAL] Sample result payload keys: {list(r0.get('payload', {}).keys())}"
+                            )
                             p = r0.get("payload", {})
-                            print(f"[EVAL] Sample payload doc_id={p.get('doc_id')}, union_doc_id={p.get('union_doc_id')}, source_doc_id={p.get('source_doc_id')}")
+                            print(
+                                f"[EVAL] Sample payload doc_id={p.get('doc_id')}, union_doc_id={p.get('union_doc_id')}, source_doc_id={p.get('source_doc_id')}"
+                            )
                         has_match = any(rid in rels for rid in ranking[:10])
                         print(f"[EVAL] Any match in top-10? {has_match}")
-                    
+
                     metrics_collectors["ndcg@5"].append(ndcg_at_k(ranking, rels, k=5))
                     metrics_collectors["ndcg@10"].append(ndcg_at_k(ranking, rels, k=10))
                     metrics_collectors["recall@5"].append(recall_at_k(ranking, rels, k=5))
                     metrics_collectors["recall@10"].append(recall_at_k(ranking, rels, k=10))
                     metrics_collectors["mrr@5"].append(mrr_at_k(ranking, rels, k=5))
                     metrics_collectors["mrr@10"].append(mrr_at_k(ranking, rels, k=10))
-                    
+
                     progress = (i + 1) / total_queries
                     progress_bar.progress(progress)
-                    eval_status.info(f"🎯 {i+1}/{total_queries} ({int(progress*100)}%) — latency: {np.mean(latencies):.0f}ms")
-                    
+                    eval_status.info(
+                        f"🎯 {i+1}/{total_queries} ({int(progress*100)}%) — latency: {np.mean(latencies):.0f}ms"
+                    )
+
                     log_interval = max(10, total_queries // 10)
                     if (i + 1) % log_interval == 0 and i > 0:
                         cur_ndcg = np.mean(metrics_collectors["ndcg@10"])
                         cur_lat = np.mean(latencies[1:]) if len(latencies) > 1 else latencies[0]
-                        log_lines.append(f"[{i+1}/{total_queries}] NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms")
+                        log_lines.append(
+                            f"[{i+1}/{total_queries}] NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms"
+                        )
                         log_area.code("\n".join(log_lines[-10:]), language="text")
-                        print(f"[EVAL] Progress {i+1}/{total_queries}: NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms")
-                
+                        print(
+                            f"[EVAL] Progress {i+1}/{total_queries}: NDCG@10={cur_ndcg:.4f}, lat={cur_lat:.0f}ms"
+                        )
+
                 progress_bar.progress(1.0)
                 eval_status.success(f"✅ Evaluation complete! ({total_queries} queries)")
-        
+
         with results_container:
             st.markdown("##### 📊 Results")
-            
+
             p95_latency = float(np.percentile(latencies, 95))
             eval_time_s = sum(latencies) / 1000
             qps = total_queries / eval_time_s if eval_time_s > 0 else 0
-            
+
             final_metrics = {
                 "ndcg@5": float(np.mean(metrics_collectors["ndcg@5"])),
                 "ndcg@10": float(np.mean(metrics_collectors["ndcg@10"])),
@@ -440,7 +487,7 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 "eval_time_s": eval_time_s,
                 "num_queries": total_queries,
             }
-            
+
             print("=" * 60)
             print("[EVAL] FINAL RESULTS:")
             print(f"[EVAL]   NDCG@5:     {final_metrics['ndcg@5']:.4f}")
@@ -454,7 +501,7 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
             print(f"[EVAL]   QPS:        {final_metrics['qps']:.2f}")
             print(f"[EVAL]   Queries:    {final_metrics['num_queries']}")
             print("=" * 60)
-            
+
             st.markdown("**Retrieval Metrics**")
             c1, c2, c3 = st.columns(3)
             with c1:
@@ -466,17 +513,17 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
             with c3:
                 st.metric("MRR@5", f"{final_metrics['mrr@5']:.4f}")
                 st.metric("MRR@10", f"{final_metrics['mrr@10']:.4f}")
-            
+
             st.markdown("**Performance**")
             c4, c5, c6, c7 = st.columns(4)
             c4.metric("Avg Latency", f"{final_metrics['avg_latency_ms']:.0f}ms")
             c5.metric("P95 Latency", f"{final_metrics['p95_latency_ms']:.0f}ms")
             c6.metric("QPS", f"{final_metrics['qps']:.2f}")
             c7.metric("Eval Time", f"{final_metrics['eval_time_s']:.1f}s")
-            
+
             with st.expander("📋 Full Results JSON"):
                 st.json(final_metrics)
-            
+
             detailed_report = {
                 "generated_at": datetime.now().isoformat(),
                 "config": {
@@ -506,17 +553,17 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                     "num_queries": final_metrics["num_queries"],
                 },
             }
-            
+
             if mode == "two_stage":
                 detailed_report["config"]["stage1_mode"] = stage1_mode
                 detailed_report["config"]["prefetch_k"] = prefetch_k
             elif mode == "three_stage":
                 detailed_report["config"]["stage1_k"] = stage1_k
                 detailed_report["config"]["stage2_k"] = stage2_k
-            
+
             if evaluation_scope == "per_dataset" and metrics_by_dataset:
                 detailed_report["metrics_by_dataset"] = metrics_by_dataset
-                
+
                 st.markdown("**Per-Dataset Results**")
                 for ds_name, ds_metrics in metrics_by_dataset.items():
                     ds_short = ds_name.split("/")[-1]
@@ -526,11 +573,11 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                         dc2.metric("Recall@10", f"{ds_metrics['recall@10']:.4f}")
                         dc3.metric("MRR@10", f"{ds_metrics['mrr@10']:.4f}")
                         dc4.metric("Latency", f"{ds_metrics['avg_latency_ms']:.0f}ms")
-            
+
             report_json = json.dumps(detailed_report, indent=2)
             timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             filename = f"eval_report__{collection}__{mode}__{timestamp}.json"
-            
+
             st.download_button(
                 label="📥 Download Detailed Report",
                 data=report_json,
@@ -538,25 +585,31 @@ def run_evaluation_with_ui(config: Dict[str, Any]):
                 mime="application/json",
                 use_container_width=True,
             )
-            
+
             st.session_state["last_eval_metrics"] = final_metrics
-        
+
     except Exception as e:
         error_msg = str(e)
-        
+
         if "not configured in this collection" in error_msg:
-            vector_name = error_msg.split("name ")[-1].split(" is")[0] if "name " in error_msg else "unknown"
-            st.error(f"❌ **Collection Mismatch**: Vector `{vector_name}` not found in collection `{collection}`")
-            st.warning(f"""
+            vector_name = (
+                error_msg.split("name ")[-1].split(" is")[0] if "name " in error_msg else "unknown"
+            )
+            st.error(
+                f"❌ **Collection Mismatch**: Vector `{vector_name}` not found in collection `{collection}`"
+            )
+            st.warning(
+                f"""
 **The selected mode `{mode}` requires vectors that don't exist in this collection.**
 
 **Suggestions:**
 - Try `single_full` mode (works with basic collections)
 - Use a collection indexed with the Visual RAG Toolkit
 - Check that the collection has the required vector types for `{mode}` mode
-            """)
+            """
+            )
         else:
             st.error(f"❌ Error: {e}")
-        
+
         with st.expander("🔍 Full Error Details"):
             st.code(traceback.format_exc(), language="text")

demo/indexing.py

@@ -1,22 +1,18 @@
 """Indexing runner with UI updates."""
 
 import hashlib
-import json
-import time
 import traceback
-from datetime import datetime
 from typing import Any, Dict, Optional
 
 import numpy as np
 import streamlit as st
 import torch
 
+from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+from demo.qdrant_utils import get_qdrant_credentials
 from visual_rag import VisualEmbedder
 from visual_rag.embedding.pooling import tile_level_mean_pooling
 from visual_rag.indexing.qdrant_indexer import QdrantIndexer
-from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
-from demo.qdrant_utils import get_qdrant_credentials
-
 
 TORCH_DTYPE_MAP = {
     "float16": torch.float16,
@@ -63,9 +59,7 @@ def run_indexing_with_ui(config: Dict[str, Any]):
 
     print(f"[INDEX] Config: collection={collection}, model={model}")
     print(f"[INDEX] Datasets: {datasets}")
-    print(
-        f"[INDEX] max_docs={max_docs}, batch_size={batch_size}, recreate={recreate}"
-    )
+    print(f"[INDEX] max_docs={max_docs}, batch_size={batch_size}, recreate={recreate}")
     print(
         f"[INDEX] torch_dtype={torch_dtype}, qdrant_dtype={qdrant_vector_dtype}, grpc={prefer_grpc}"
     )
@@ -83,9 +77,7 @@ def run_indexing_with_ui(config: Dict[str, Any]):
 
             print(f"[INDEX] Loading embedder: {model}")
             torch_dtype_obj = TORCH_DTYPE_MAP.get(torch_dtype, torch.float16)
-            output_dtype_obj = (
-                np.float16 if qdrant_vector_dtype == "float16" else np.float32
-            )
+            output_dtype_obj = np.float16 if qdrant_vector_dtype == "float16" else np.float32
             embedder = VisualEmbedder(
                 model_name=model,
                 torch_dtype=torch_dtype_obj,
@@ -140,9 +132,7 @@ def run_indexing_with_ui(config: Dict[str, Any]):
                 ds_container = st.container()
 
                 with ds_container:
-                    st.markdown(
-                        f"**Dataset {ds_idx + 1}/{len(datasets)}: `{ds_short}`**"
-                    )
+                    st.markdown(f"**Dataset {ds_idx + 1}/{len(datasets)}: `{ds_short}`**")
 
                     load_status = st.empty()
                     load_status.info(f"Loading dataset `{ds_short}`...")
@@ -172,9 +162,7 @@ def run_indexing_with_ui(config: Dict[str, Any]):
                                 failed += 1
                                 continue
 
-                            status_text.text(
-                                f"Processing {i + 1}/{total}: {doc_id[:30]}..."
-                            )
+                            status_text.text(f"Processing {i + 1}/{total}: {doc_id[:30]}...")
 
                             embeddings, token_infos = embedder.embed_images(
                                 [image],

demo/qdrant_utils.py

@@ -9,7 +9,7 @@ import streamlit as st
 
 def get_qdrant_credentials() -> Tuple[Optional[str], Optional[str]]:
     """Get Qdrant credentials from session state or environment variables.
-    
+
     Priority: session_state > QDRANT_URL/QDRANT_API_KEY > legacy env vars
     """
     url = (
@@ -28,6 +28,7 @@ def get_qdrant_credentials() -> Tuple[Optional[str], Optional[str]]:
 def init_qdrant_client_with_creds(url: str, api_key: str):
     try:
         from qdrant_client import QdrantClient
+
         if not url:
             return None, "QDRANT_URL not configured"
         client = QdrantClient(url=url, api_key=api_key, timeout=60)
@@ -47,6 +48,7 @@ def init_qdrant_client():
 def init_embedder(model_name: str):
     try:
         from visual_rag import VisualEmbedder
+
         return VisualEmbedder(model_name=model_name), None
     except Exception as e:
         return None, f"{e}\n\n{traceback.format_exc()}"
@@ -72,7 +74,9 @@ def get_collection_stats(collection_name: str) -> Dict[str, Any]:
         return {"error": err}
     try:
         info = client.get_collection(collection_name)
-        vectors_config = getattr(getattr(getattr(info, "config", None), "params", None), "vectors", None)
+        vectors_config = getattr(
+            getattr(getattr(info, "config", None), "params", None), "vectors", None
+        )
         vector_info = {}
         if vectors_config is not None:
             if hasattr(vectors_config, "items"):
@@ -96,7 +100,9 @@ def get_collection_stats(collection_name: str) -> Dict[str, Any]:
                     }
             elif hasattr(vectors_config, "size"):
                 on_disk = getattr(vectors_config, "on_disk", None)
-                datatype = str(getattr(vectors_config, "datatype", "Float32")).replace("Datatype.", "")
+                datatype = str(getattr(vectors_config, "datatype", "Float32")).replace(
+                    "Datatype.", ""
+                )
                 multivec = getattr(vectors_config, "multivector_config", None)
                 vector_info["default"] = {
                     "size": getattr(vectors_config, "size", None),
@@ -117,12 +123,15 @@ def get_collection_stats(collection_name: str) -> Dict[str, Any]:
 
 
 @st.cache_data(ttl=60)
-def sample_points_cached(collection_name: str, n: int, seed: int, _url: str, _api_key: str) -> List[Dict[str, Any]]:
+def sample_points_cached(
+    collection_name: str, n: int, seed: int, _url: str, _api_key: str
+) -> List[Dict[str, Any]]:
     client, err = init_qdrant_client_with_creds(_url, _api_key)
     if client is None:
         return []
     try:
         import random
+
         rng = random.Random(seed)
         points, _ = client.scroll(
             collection_name=collection_name,
@@ -136,10 +145,12 @@ def sample_points_cached(collection_name: str, n: int, seed: int, _url: str, _ap
         results = []
         for p in sampled:
             payload = dict(p.payload) if p.payload else {}
-            results.append({
-                "id": str(p.id),
-                "payload": payload,
-            })
+            results.append(
+                {
+                    "id": str(p.id),
+                    "payload": payload,
+                }
+            )
         return results
     except Exception:
         return []
@@ -181,14 +192,16 @@ def search_collection(
     top_k: int = 10,
     mode: str = "single_full",
     prefetch_k: int = 256,
-    stage1_mode: str = "tokens_vs_tiles",
+    stage1_mode: str = "tokens_vs_standard_pooling",
     stage1_k: int = 1000,
     stage2_k: int = 300,
     model_name: str = "vidore/colSmol-500M",
 ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
     try:
         import traceback
+
         from visual_rag.retrieval import MultiVectorRetriever
+
         retriever = MultiVectorRetriever(
             collection_name=collection_name,
             model_name=model_name,
@@ -215,4 +228,5 @@ def search_collection(
         return results, None
     except Exception as e:
         import traceback
+
         return [], f"{e}\n\n{traceback.format_exc()}"

demo/test_qdrant_connection.py

@@ -7,27 +7,29 @@ from pathlib import Path
 
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from dotenv import load_dotenv
+from dotenv import load_dotenv  # noqa: E402
+
 load_dotenv(Path(__file__).parent.parent / ".env")
 load_dotenv(Path(__file__).parent.parent.parent / ".env")
 
+
 def test_connection():
     from qdrant_client import QdrantClient
     from qdrant_client.http import models
-    
+
     url = os.getenv("QDRANT_URL")
     api_key = os.getenv("QDRANT_API_KEY")
-    
+
     print(f"URL: {url}")
     print(f"API Key: {'***' + api_key[-4:] if api_key else 'NOT SET'}")
-    
+
     if not url or not api_key:
         print("ERROR: QDRANT_URL or QDRANT_API_KEY not set")
         return
-    
+
     print("\n1. Creating client...")
     client = QdrantClient(url=url, api_key=api_key, timeout=60)
-    
+
     print("\n2. Getting collections...")
     try:
         collections = client.get_collections()
@@ -37,22 +39,22 @@ def test_connection():
     except Exception as e:
         print(f"   ERROR: {e}")
         return
-    
+
     test_collection = "_test_visual_rag_toolkit"
-    
+
     print(f"\n3. Checking if '{test_collection}' exists...")
     exists = any(c.name == test_collection for c in collections.collections)
     print(f"   Exists: {exists}")
-    
+
     if exists:
-        print(f"\n4. Deleting test collection...")
+        print("\n4. Deleting test collection...")
         try:
             client.delete_collection(test_collection)
             print("   Deleted")
         except Exception as e:
             print(f"   ERROR: {e}")
-    
-    print(f"\n5. Creating SIMPLE collection (single vector)...")
+
+    print("\n5. Creating SIMPLE collection (single vector)...")
     try:
         client.create_collection(
             collection_name=test_collection,
@@ -67,15 +69,15 @@ def test_connection():
         print("\n   This means basic collection creation is failing.")
         print("   Check your Qdrant Cloud cluster status/limits.")
         return
-    
-    print(f"\n6. Deleting test collection...")
+
+    print("\n6. Deleting test collection...")
     try:
         client.delete_collection(test_collection)
         print("   Deleted")
     except Exception as e:
         print(f"   ERROR: {e}")
-    
-    print(f"\n7. Creating MULTI-VECTOR collection (like visual-rag)...")
+
+    print("\n7. Creating MULTI-VECTOR collection (like visual-rag)...")
     try:
         client.create_collection(
             collection_name=test_collection,
@@ -102,17 +104,17 @@ def test_connection():
         print("\n   Multi-vector collection failed but simple worked.")
         print("   Your Qdrant version may not support multi-vector.")
         return
-    
-    print(f"\n8. Final cleanup...")
+
+    print("\n8. Final cleanup...")
     try:
         client.delete_collection(test_collection)
         print("   Deleted")
     except Exception as e:
         print(f"   ERROR: {e}")
-    
-    print("\n" + "="*50)
+
+    print("\n" + "=" * 50)
     print("ALL TESTS PASSED - Qdrant connection is working!")
-    print("="*50)
+    print("=" * 50)
 
 
 if __name__ == "__main__":

demo/ui/__init__.py

@@ -1,10 +1,10 @@
 """UI components for the demo app."""
 
+from demo.ui.benchmark import render_benchmark_tab
 from demo.ui.header import render_header
+from demo.ui.playground import render_playground_tab
 from demo.ui.sidebar import render_sidebar
 from demo.ui.upload import render_upload_tab
-from demo.ui.playground import render_playground_tab
-from demo.ui.benchmark import render_benchmark_tab
 
 __all__ = [
     "render_header",

demo/ui/benchmark.py

@@ -7,6 +7,12 @@ import altair as alt
 import pandas as pd
 import streamlit as st
 
+from demo.commands import (
+    build_eval_command,
+    build_index_command,
+    generate_python_eval_code,
+    generate_python_index_code,
+)
 from demo.config import (
     AVAILABLE_MODELS,
     BENCHMARK_DATASETS,
@@ -14,47 +20,52 @@ from demo.config import (
     RETRIEVAL_MODES,
     STAGE1_MODES,
 )
-from demo.qdrant_utils import get_qdrant_credentials, get_collections
-from demo.commands import build_index_command, build_eval_command, generate_python_eval_code, generate_python_index_code
-from demo.results import get_available_results, load_results_file
 from demo.evaluation import run_evaluation_with_ui
 from demo.indexing import run_indexing_with_ui
+from demo.qdrant_utils import get_collections, get_qdrant_credentials
+from demo.results import get_available_results, load_results_file
 
 
 def render_benchmark_tab():
     st.subheader("📊 Benchmarking")
-    
+
     tab_index, tab_eval, tab_results = st.tabs(["Indexing", "Evaluation", "Results"])
-    
+
     url, api_key = get_qdrant_credentials()
     collections = get_collections(url, api_key)
-    
+
     with tab_index:
         render_benchmark_indexing(collections)
-    
+
     with tab_eval:
         render_benchmark_evaluation(collections)
-    
+
     with tab_results:
         render_benchmark_results()
 
 
 def render_benchmark_indexing(collections: List[str]):
     st.caption("Create a new collection with benchmark datasets")
-    
+
     c1, c2, c3 = st.columns(3)
     with c1:
-        datasets = st.multiselect("Datasets", BENCHMARK_DATASETS, default=BENCHMARK_DATASETS, key="bi_ds")
+        datasets = st.multiselect(
+            "Datasets", BENCHMARK_DATASETS, default=BENCHMARK_DATASETS, key="bi_ds"
+        )
     with c2:
         model = st.selectbox("Model", AVAILABLE_MODELS, key="bi_model")
     with c3:
         model_short = model.split("/")[-1].replace("-", "_").replace(".", "_")
-        collection = st.text_input("New Collection Name", value=f"vidore_{len(datasets)}ds__{model_short}", key="bi_coll")
-    
+        collection = st.text_input(
+            "New Collection Name", value=f"vidore_{len(datasets)}ds__{model_short}", key="bi_coll"
+        )
+
     sel_docs = sum(DATASET_STATS.get(d, {}).get("docs", 0) for d in datasets)
     sel_queries = sum(DATASET_STATS.get(d, {}).get("queries", 0) for d in datasets)
-    st.markdown(f"🎯 **Selected:** {len(datasets)} dataset(s) — **{sel_docs:,}** docs, **{sel_queries:,}** queries")
-    
+    st.markdown(
+        f"🎯 **Selected:** {len(datasets)} dataset(s) — **{sel_docs:,}** docs, **{sel_queries:,}** queries"
+    )
+
     c4, c5, c6, c7 = st.columns(4)
     with c4:
         crop = st.toggle("Crop", value=True, key="bi_crop")
@@ -64,11 +75,11 @@ def render_benchmark_indexing(collections: List[str]):
         grpc = st.toggle("gRPC", value=True, key="bi_grpc")
     with c7:
         recreate = st.toggle("Recreate", value=False, key="bi_recreate")
-    
+
     crop_pct = st.slider("Crop %", 0.8, 0.99, 0.99, 0.01, key="bi_crop_pct") if crop else 0.99
-    
+
     st.markdown("---")
-    
+
     col_max, col_batch, col_torch, col_qdrant = st.columns([2, 2, 1, 1])
     with col_max:
         max_docs_val = max(sel_docs, 1)
@@ -78,30 +89,47 @@ def render_benchmark_indexing(collections: List[str]):
             max_value=max_docs_val,
             value=max_docs_val,
             key="bi_max_docs",
-            help="Limit docs per dataset. Useful for quick tests."
+            help="Limit docs per dataset. Useful for quick tests.",
         )
     with col_batch:
-        batch_size = st.number_input("Batch Size", min_value=1, max_value=16, value=4, key="bi_batch")
+        batch_size = st.number_input(
+            "Batch Size", min_value=1, max_value=16, value=4, key="bi_batch"
+        )
     with col_torch:
-        torch_dtype = st.selectbox("Torch dtype", ["float16", "float32"], index=0, key="bi_torch_dtype")
+        torch_dtype = st.selectbox(
+            "Torch dtype", ["float16", "float32"], index=0, key="bi_torch_dtype"
+        )
     with col_qdrant:
-        qdrant_dtype = st.selectbox("Qdrant dtype", ["float16", "float32"], index=0, key="bi_qdrant_dtype")
-    
-    effective_docs = min(max_docs * len(datasets), sel_docs) if max_docs < max_docs_val else sel_docs
-    
+        qdrant_dtype = st.selectbox(
+            "Qdrant dtype", ["float16", "float32"], index=0, key="bi_qdrant_dtype"
+        )
+
+    effective_docs = (
+        min(max_docs * len(datasets), sel_docs) if max_docs < max_docs_val else sel_docs
+    )
+
     config = {
-        "datasets": datasets, "model": model, "collection": collection,
-        "crop_empty": crop, "crop_percentage": crop_pct,
-        "no_cloudinary": not cloudinary, "recreate": recreate, "resume": False,
-        "prefer_grpc": grpc, "batch_size": batch_size, "upload_batch_size": 8,
-        "qdrant_timeout": 180, "qdrant_retries": 5,
-        "torch_dtype": torch_dtype, "qdrant_vector_dtype": qdrant_dtype,
+        "datasets": datasets,
+        "model": model,
+        "collection": collection,
+        "crop_empty": crop,
+        "crop_percentage": crop_pct,
+        "no_cloudinary": not cloudinary,
+        "recreate": recreate,
+        "resume": False,
+        "prefer_grpc": grpc,
+        "batch_size": batch_size,
+        "upload_batch_size": 8,
+        "qdrant_timeout": 180,
+        "qdrant_retries": 5,
+        "torch_dtype": torch_dtype,
+        "qdrant_vector_dtype": qdrant_dtype,
         "max_docs": max_docs if max_docs < max_docs_val else None,
     }
-    
+
     cmd = build_index_command(config)
     python_code = generate_python_index_code(config)
-    
+
     col_cmd, col_info = st.columns([2, 1])
     with col_cmd:
         code_tab1, code_tab2 = st.tabs(["🐚 Bash", "🐍 Python"])
@@ -111,14 +139,16 @@ def render_benchmark_indexing(collections: List[str]):
             st.code(python_code, language="python")
     with col_info:
         st.markdown("<br><br><br>", unsafe_allow_html=True)
-        
+
         st.metric("Docs to Index", f"{effective_docs:,}")
         st.metric("Model", model.split("/")[-1])
         if effective_docs < sel_docs:
             st.caption(f"Limited from {sel_docs:,} total")
         st.divider()
-        run_index = st.button("🚀 Run Index", type="primary", key="bi_run", use_container_width=True)
-    
+        run_index = st.button(
+            "🚀 Run Index", type="primary", key="bi_run", use_container_width=True
+        )
+
     if run_index:
         if not collection:
             st.error("Please provide a collection name")
@@ -130,38 +160,42 @@ def render_benchmark_indexing(collections: List[str]):
 
 def render_benchmark_evaluation(collections: List[str]):
     collection = st.session_state.get("active_collection")
-    
+
     if not collection:
         st.warning("⚠️ Select a collection from the sidebar first")
         return
-    
+
     st.info(f"**Collection:** `{collection}` (from sidebar)")
-    
+
     all_docs = sum(DATASET_STATS.get(d, {}).get("docs", 0) for d in BENCHMARK_DATASETS)
     all_queries = sum(DATASET_STATS.get(d, {}).get("queries", 0) for d in BENCHMARK_DATASETS)
-    st.markdown(f"📊 **Available:** {len(BENCHMARK_DATASETS)} datasets — **{all_docs:,}** docs, **{all_queries:,}** queries")
-    
+    st.markdown(
+        f"📊 **Available:** {len(BENCHMARK_DATASETS)} datasets — **{all_docs:,}** docs, **{all_queries:,}** queries"
+    )
+
     c1, c2 = st.columns([3, 1])
     with c1:
         st.multiselect("Datasets", BENCHMARK_DATASETS, default=BENCHMARK_DATASETS, key="be_ds")
     with c2:
         model = st.selectbox("Model", AVAILABLE_MODELS, key="be_model")
-    
+
     datasets = st.session_state.get("be_ds", BENCHMARK_DATASETS)
     sel_docs = sum(DATASET_STATS.get(d, {}).get("docs", 0) for d in datasets)
     sel_queries = sum(DATASET_STATS.get(d, {}).get("queries", 0) for d in datasets)
-    st.markdown(f"🎯 **Selected:** {len(datasets)} dataset(s) — **{sel_docs:,}** docs, **{sel_queries:,}** queries")
-    
+    st.markdown(
+        f"🎯 **Selected:** {len(datasets)} dataset(s) — **{sel_docs:,}** docs, **{sel_queries:,}** queries"
+    )
+
     st.markdown("---")
-    
+
     col_mode, col_topk = st.columns([2, 1])
     with col_mode:
         mode = st.selectbox("Mode", RETRIEVAL_MODES, key="be_mode")
     with col_topk:
         top_k = st.slider("Top K", 10, 100, 100, key="be_topk")
-    
-    stage1_mode, prefetch_k, stage1_k, stage2_k = "tokens_vs_tiles", 256, 1000, 300
-    
+
+    stage1_mode, prefetch_k, stage1_k, stage2_k = "tokens_vs_standard_pooling", 256, 1000, 300
+
     if mode == "two_stage":
         cc1, cc2 = st.columns(2)
         with cc1:
@@ -174,9 +208,9 @@ def render_benchmark_evaluation(collections: List[str]):
             stage1_k = st.number_input("Stage1 K", 100, 5000, 1000, key="be_s1k")
         with cc2:
             stage2_k = st.number_input("Stage2 K", 50, 1000, 300, key="be_s2k")
-    
+
     st.markdown("---")
-    
+
     col_scope, _, col_grpc, col_nq = st.columns([2, 0.5, 1, 2])
     with col_scope:
         scope = st.selectbox("Scope", ["union", "per_dataset"], key="be_scope")
@@ -187,33 +221,39 @@ def render_benchmark_evaluation(collections: List[str]):
     with col_nq:
         max_q_val = max(sel_queries, 1)
         max_queries = st.number_input(
-            "Max Queries", 
-            min_value=1, 
-            max_value=max_q_val, 
-            value=max_q_val, 
+            "Max Queries",
+            min_value=1,
+            max_value=max_q_val,
+            value=max_q_val,
             key="be_max_queries",
-            help="Limit number of queries to evaluate (useful for quick tests)"
+            help="Limit number of queries to evaluate (useful for quick tests)",
         )
-    
+
     result_prefix_val = st.session_state.get("be_prefix", "")
-    
+
     config = {
-        "datasets": datasets, "model": model, "collection": collection,
-        "mode": mode, "top_k": top_k, "evaluation_scope": scope,
+        "datasets": datasets,
+        "model": model,
+        "collection": collection,
+        "mode": mode,
+        "top_k": top_k,
+        "evaluation_scope": scope,
         "prefer_grpc": grpc,
         "torch_dtype": "float16",
         "qdrant_vector_dtype": "float16",
         "qdrant_timeout": 180,
-        "stage1_mode": stage1_mode, "prefetch_k": prefetch_k,
-        "stage1_k": stage1_k, "stage2_k": stage2_k,
+        "stage1_mode": stage1_mode,
+        "prefetch_k": prefetch_k,
+        "stage1_k": stage1_k,
+        "stage2_k": stage2_k,
         "result_prefix": result_prefix_val,
         "max_queries": max_queries,
     }
-    
+
     cmd = build_eval_command(config)
-    
+
     python_code = generate_python_eval_code(config)
-    
+
     col_cmd, col_info = st.columns([2, 1])
     with col_cmd:
         code_tab1, code_tab2 = st.tabs(["🐚 Bash", "🐍 Python"])
@@ -223,7 +263,7 @@ def render_benchmark_evaluation(collections: List[str]):
             st.code(python_code, language="python")
     with col_info:
         st.markdown("<br><br><br>", unsafe_allow_html=True)
-        
+
         mode_desc = {
             "single_full": "🔹 **Single Full**: Query all visual tokens against full document embeddings in one pass.",
             "single_tiles": "🔸 **Single Tiles**: Query against tile-level embeddings only.",
@@ -239,9 +279,9 @@ def render_benchmark_evaluation(collections: List[str]):
         st.markdown(scope_desc.get(scope, ""))
         st.divider()
         st.text_input("Result Prefix", placeholder="optional prefix for output", key="be_prefix")
-        
+
         run_eval = st.button("🚀 Run Eval", type="primary", key="be_run", use_container_width=True)
-    
+
     if run_eval:
         if not collection:
             st.error("Please select a collection first")
@@ -251,19 +291,19 @@ def render_benchmark_evaluation(collections: List[str]):
 
 def render_benchmark_results():
     st.markdown("##### Load Results")
-    
+
     available = get_available_results()
-    
+
     if not available:
         st.info("No results found")
         return
-    
+
     default_select = []
     if st.session_state.get("auto_select_result"):
         auto = st.session_state.pop("auto_select_result")
         if auto in [str(p) for p in available]:
             default_select = [auto]
-    
+
     selected = st.multiselect(
         "Result files",
         options=[str(p) for p in available],
@@ -271,7 +311,7 @@ def render_benchmark_results():
         default=default_select,
         key="br_files",
     )
-    
+
     for path in selected:
         data = load_results_file(Path(path))
         if data:
@@ -285,71 +325,108 @@ def render_result_card(data: Dict[str, Any], filename: str):
         c2.metric("Mode", data.get("mode", "?"))
         c3.metric("Top K", data.get("top_k", "?"))
         c4.metric("Time", f"{data.get('eval_wall_time_s', 0):.0f}s")
-        
+
         metrics = data.get("metrics_by_dataset", {})
         if not metrics:
             st.warning("No metrics data")
             return
-        
+
         rows = []
         for ds, m in metrics.items():
-            rows.append({
-                "Dataset": ds.split("/")[-1].replace("_v2", ""),
-                "NDCG@5": m.get("ndcg@5", 0),
-                "NDCG@10": m.get("ndcg@10", 0),
-                "Recall@5": m.get("recall@5", 0),
-                "Recall@10": m.get("recall@10", 0),
-                "MRR@10": m.get("mrr@10", 0),
-                "Latency": m.get("avg_latency_ms", 0),
-                "QPS": m.get("qps", 0),
-            })
-        
+            rows.append(
+                {
+                    "Dataset": ds.split("/")[-1].replace("_v2", ""),
+                    "NDCG@5": m.get("ndcg@5", 0),
+                    "NDCG@10": m.get("ndcg@10", 0),
+                    "Recall@5": m.get("recall@5", 0),
+                    "Recall@10": m.get("recall@10", 0),
+                    "MRR@10": m.get("mrr@10", 0),
+                    "Latency": m.get("avg_latency_ms", 0),
+                    "QPS": m.get("qps", 0),
+                }
+            )
+
         df = pd.DataFrame(rows)
-        
+
         st.dataframe(
-            df.style.format({
-                "NDCG@5": "{:.4f}", "NDCG@10": "{:.4f}",
-                "Recall@5": "{:.4f}", "Recall@10": "{:.4f}",
-                "MRR@10": "{:.4f}", "Latency": "{:.1f}", "QPS": "{:.2f}"
-            }),
-            hide_index=True, use_container_width=True
+            df.style.format(
+                {
+                    "NDCG@5": "{:.4f}",
+                    "NDCG@10": "{:.4f}",
+                    "Recall@5": "{:.4f}",
+                    "Recall@10": "{:.4f}",
+                    "MRR@10": "{:.4f}",
+                    "Latency": "{:.1f}",
+                    "QPS": "{:.2f}",
+                }
+            ),
+            hide_index=True,
+            use_container_width=True,
         )
-        
+
         chart_data = []
         for ds, m in metrics.items():
             ds_short = ds.split("/")[-1].replace("_v2", "").replace("_", " ").title()
-            chart_data.append({"Dataset": ds_short, "Metric": "NDCG@10", "Value": m.get("ndcg@10", 0)})
-            chart_data.append({"Dataset": ds_short, "Metric": "Recall@10", "Value": m.get("recall@10", 0)})
-            chart_data.append({"Dataset": ds_short, "Metric": "MRR@10", "Value": m.get("mrr@10", 0)})
-        
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "NDCG@10", "Value": m.get("ndcg@10", 0)}
+            )
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "Recall@10", "Value": m.get("recall@10", 0)}
+            )
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "MRR@10", "Value": m.get("mrr@10", 0)}
+            )
+
         chart_df = pd.DataFrame(chart_data)
-        
-        chart = alt.Chart(chart_df).mark_bar().encode(
-            x=alt.X("Dataset:N", title=None),
-            y=alt.Y("Value:Q", scale=alt.Scale(domain=[0, 1]), title="Score"),
-            color=alt.Color("Metric:N", scale=alt.Scale(scheme="tableau10")),
-            xOffset="Metric:N",
-            tooltip=["Dataset", "Metric", alt.Tooltip("Value:Q", format=".4f")]
-        ).properties(height=300, title="Metrics by Dataset")
-        
+
+        chart = (
+            alt.Chart(chart_df)
+            .mark_bar()
+            .encode(
+                x=alt.X("Dataset:N", title=None),
+                y=alt.Y("Value:Q", scale=alt.Scale(domain=[0, 1]), title="Score"),
+                color=alt.Color("Metric:N", scale=alt.Scale(scheme="tableau10")),
+                xOffset="Metric:N",
+                tooltip=["Dataset", "Metric", alt.Tooltip("Value:Q", format=".4f")],
+            )
+            .properties(height=300, title="Metrics by Dataset")
+        )
+
         st.altair_chart(chart, use_container_width=True)
-        
-        latency_data = [{"Dataset": ds.split("/")[-1].replace("_v2", ""), "Latency (ms)": m.get("avg_latency_ms", 0), "QPS": m.get("qps", 0)} for ds, m in metrics.items()]
+
+        latency_data = [
+            {
+                "Dataset": ds.split("/")[-1].replace("_v2", ""),
+                "Latency (ms)": m.get("avg_latency_ms", 0),
+                "QPS": m.get("qps", 0),
+            }
+            for ds, m in metrics.items()
+        ]
         latency_df = pd.DataFrame(latency_data)
-        
+
         c1, c2 = st.columns(2)
         with c1:
-            lat_chart = alt.Chart(latency_df).mark_bar(color="#ff6b6b").encode(
-                x=alt.X("Dataset:N"),
-                y=alt.Y("Latency (ms):Q"),
-                tooltip=["Dataset", alt.Tooltip("Latency (ms):Q", format=".1f")]
-            ).properties(height=200, title="Avg Latency")
+            lat_chart = (
+                alt.Chart(latency_df)
+                .mark_bar(color="#ff6b6b")
+                .encode(
+                    x=alt.X("Dataset:N"),
+                    y=alt.Y("Latency (ms):Q"),
+                    tooltip=["Dataset", alt.Tooltip("Latency (ms):Q", format=".1f")],
+                )
+                .properties(height=200, title="Avg Latency")
+            )
             st.altair_chart(lat_chart, use_container_width=True)
-        
+
         with c2:
-            qps_chart = alt.Chart(latency_df).mark_bar(color="#4ecdc4").encode(
-                x=alt.X("Dataset:N"),
-                y=alt.Y("QPS:Q"),
-                tooltip=["Dataset", alt.Tooltip("QPS:Q", format=".2f")]
-            ).properties(height=200, title="QPS (Queries/sec)")
+            qps_chart = (
+                alt.Chart(latency_df)
+                .mark_bar(color="#4ecdc4")
+                .encode(
+                    x=alt.X("Dataset:N"),
+                    y=alt.Y("QPS:Q"),
+                    tooltip=["Dataset", alt.Tooltip("QPS:Q", format=".2f")],
+                )
+                .properties(height=200, title="QPS (Queries/sec)")
+            )
             st.altair_chart(qps_chart, use_container_width=True)

demo/ui/header.py

@@ -4,7 +4,8 @@ import streamlit as st
 
 
 def render_header():
-    st.markdown("""
+    st.markdown(
+        """
     <div style="text-align: center; padding: 10px 0 15px 0;">
         <h1 style="
             font-family: 'Georgia', serif;
@@ -35,4 +36,6 @@ def render_header():
             </a>
         </p>
     </div>
-    """, unsafe_allow_html=True)
+    """,
+        unsafe_allow_html=True,
+    )

demo/ui/playground.py

@@ -4,8 +4,8 @@ import streamlit as st
 
 from demo.config import AVAILABLE_MODELS, RETRIEVAL_MODES, STAGE1_MODES
 from demo.qdrant_utils import (
-    get_qdrant_credentials,
     get_collections,
+    get_qdrant_credentials,
     sample_points_cached,
     search_collection,
 )
@@ -14,32 +14,32 @@ from visual_rag.retrieval import MultiVectorRetriever
 
 def render_playground_tab():
     st.subheader("🎮 Playground")
-    
+
     active_collection = st.session_state.get("active_collection")
     url, api_key = get_qdrant_credentials()
-    
+
     if not active_collection:
         collections = get_collections(url, api_key)
         if collections:
             active_collection = collections[0]
-    
+
     if not active_collection:
         st.warning("No collection available. Upload documents or select a collection.")
         return
-    
+
     points_for_model = sample_points_cached(active_collection, 1, 0, url, api_key)
     model_name = None
     if points_for_model:
         model_name = points_for_model[0].get("payload", {}).get("model_name")
     if not model_name:
         model_name = AVAILABLE_MODELS[1]
-    
+
     model_short = model_name.split("/")[-1] if model_name else "unknown"
     cache_key = f"{active_collection}_{model_name}"
-    
+
     if st.session_state.get("loaded_model_key") != cache_key:
         st.session_state["model_loaded"] = False
-    
+
     col_info, col_model = st.columns([2, 1])
     with col_info:
         st.info(f"**Collection:** `{active_collection}`")
@@ -47,21 +47,26 @@ def render_playground_tab():
         if not st.session_state.get("model_loaded"):
             with st.spinner(f"Loading {model_short}..."):
                 try:
-                    _ = MultiVectorRetriever(collection_name=active_collection, model_name=model_name)
+                    _ = MultiVectorRetriever(
+                        collection_name=active_collection, model_name=model_name
+                    )
                     st.session_state["model_loaded"] = True
                     st.session_state["loaded_model_key"] = cache_key
                     st.session_state["loaded_model_name"] = model_name
-                except Exception as e:
+                except Exception:
                     st.warning(f"Failed: {model_short}")
-        
+
         if st.session_state.get("model_loaded"):
-            st.markdown(f"✅ Found <span style='color:#e74c3c;font-weight:bold;'>{model_short}</span> model", unsafe_allow_html=True)
-    
+            st.markdown(
+                f"✅ Found <span style='color:#e74c3c;font-weight:bold;'>{model_short}</span> model",
+                unsafe_allow_html=True,
+            )
+
     with st.expander("📦 Sample Points Explorer", expanded=True):
         render_sample_explorer(active_collection, url, api_key)
-    
+
     st.divider()
-    
+
     st.subheader("🔍 RAG Query")
     render_rag_query_interface(active_collection, model_name)
 
@@ -80,15 +85,15 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
     doc_id = p.get("doc_id") or p.get("union_doc_id") or p.get("source_doc_id") or "?"
     corpus_id = p.get("corpus-id") or p.get("source_doc_id") or "?"
     dataset = p.get("dataset") or p.get("source") or None
-    model = (p.get("model_name") or p.get("model") or None)
+    model = p.get("model_name") or p.get("model") or None
     model = model.split("/")[-1] if isinstance(model, str) else None
     doc_name = p.get("doc-id") or p.get("filename") or "Unknown"
-    
+
     num_tiles = p.get("num_tiles")
     visual_tokens = p.get("index_recovery_num_visual_tokens") or p.get("num_visual_tokens")
     patches_per_tile = p.get("patches_per_tile")
     torch_dtype = p.get("torch_dtype")
-    
+
     orig_w = p.get("original_width")
     orig_h = p.get("original_height")
     crop_w = p.get("cropped_width")
@@ -97,9 +102,9 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
     resize_h = p.get("resized_height")
     crop_pct = p.get("crop_empty_percentage_to_remove")
     crop_enabled = bool(p.get("crop_empty_enabled", False))
-    
+
     col_meta, col_img = st.columns([1, 2])
-    
+
     with col_meta:
         st.markdown("##### 📄 Document Info")
         st.markdown(f"**📁 Doc:** {doc_name}")
@@ -109,7 +114,7 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
             st.markdown(f"**🔑 Doc ID:** `{str(doc_id)[:20]}...`")
         if not _is_missing(corpus_id) and str(corpus_id) != "?":
             st.markdown(f"**📋 Corpus ID:** {corpus_id}")
-        
+
         if score is not None:
             st.divider()
             st.markdown("##### 🎯 Relevance")
@@ -117,7 +122,7 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
                 st.markdown(f"**Relative:** 🟢 {rel_pct:.1f}%")
                 st.progress(rel_pct / 100)
             st.caption(f"Raw score: {score:.4f}")
-        
+
         st.divider()
         visual_rows = []
         if not _is_missing(model):
@@ -134,7 +139,7 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
             st.markdown("##### 🎨 Visual Metadata")
             for k, v in visual_rows:
                 st.markdown(f"**{k}:** {v}")
-        
+
         st.divider()
         dim_rows = []
         if not _is_missing(orig_w) and not _is_missing(orig_h):
@@ -152,33 +157,37 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
                 st.markdown(f"**Crop %:** {int(float(crop_pct) * 100)}%")
             except Exception:
                 pass
-    
+
     with col_img:
         st.markdown("##### 📷 Document Page")
         tabs = st.tabs(["🖼️ Original", "📷 Resized", "✂️ Cropped"])
-        
+
         url_o = p.get("original_url")
         url_r = p.get("resized_url") or p.get("page")
         url_c = p.get("cropped_url")
-        
+
         with tabs[0]:
             if url_o:
                 st.image(url_o, width=600)
                 st.caption(f"📐 **{orig_w}×{orig_h}**")
             else:
                 st.info("No original image available")
-        
+
         with tabs[1]:
             if url_r:
                 st.image(url_r, width=600)
                 st.caption(f"📐 **{resize_w}×{resize_h}**")
             else:
                 st.info("No resized image available")
-        
+
         with tabs[2]:
             if url_c:
                 # Display on a checkerboard background to make the crop boundary obvious.
-                w_caption = f"{crop_w}×{crop_h}" if (not _is_missing(crop_w) and not _is_missing(crop_h)) else None
+                w_caption = (
+                    f"{crop_w}×{crop_h}"
+                    if (not _is_missing(crop_w) and not _is_missing(crop_h))
+                    else None
+                )
                 pct_caption = None
                 if not _is_missing(crop_pct):
                     try:
@@ -215,7 +224,7 @@ def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: flo
                     st.caption(" | ".join(cap))
             else:
                 st.info("No cropped image available")
-        
+
         with st.expander("🔗 Image URLs"):
             if url_o:
                 st.code(url_o, language=None)
@@ -235,7 +244,7 @@ def render_sample_explorer(collection_name: str, url: str, api_key: str):
             datasets.add(ds)
         if did := (p.get("doc-id") or p.get("filename")):
             doc_ids.add(did)
-    
+
     c1, c2, c3, c4 = st.columns([1, 1, 2, 1])
     with c1:
         n_samples = st.slider("Samples", 1, 20, 3, key="pg_n")
@@ -246,28 +255,28 @@ def render_sample_explorer(collection_name: str, url: str, api_key: str):
     with c4:
         st.write("")
         do_sample = st.button("🎲 Sample", type="primary", key="pg_sample_btn")
-    
+
     if do_sample:
         points = sample_points_cached(collection_name, n_samples * 5, seed, url, api_key)
         if filter_ds != "All":
             points = [p for p in points if p.get("payload", {}).get("dataset") == filter_ds]
         points = points[:n_samples]
         st.session_state["pg_points"] = points
-    
+
     points = st.session_state.get("pg_points", [])
-    
+
     if not points:
         st.caption("Click 'Sample' to load documents")
         return
-    
+
     st.success(f"**{len(points)} points loaded**")
-    
+
     for i, pt in enumerate(points):
         p = pt.get("payload", {})
-        
+
         filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
         page_num = p.get("page_number") or p.get("page") or "?"
-        
+
         with st.expander(f"**{i+1}.** {str(filename)[:40]} - Page {page_num}", expanded=(i == 0)):
             render_document_details(pt, p)
 
@@ -275,26 +284,26 @@ def render_sample_explorer(collection_name: str, url: str, api_key: str):
 def render_rag_query_interface(collection_name: str, model_name: str = None):
     if not collection_name:
         return
-    
+
     url, api_key = get_qdrant_credentials()
-    
+
     if not model_name:
         points = sample_points_cached(collection_name, 1, 0, url, api_key)
         if points:
             model_name = points[0].get("payload", {}).get("model_name")
         if not model_name:
             model_name = AVAILABLE_MODELS[1]
-    
+
     st.caption(f"Model: **{model_name.split('/')[-1] if model_name else 'auto'}**")
-    
+
     c1, c2, c3 = st.columns([2, 1, 1])
     with c2:
         mode = st.selectbox("Mode", RETRIEVAL_MODES, index=0, key="q_mode")
     with c3:
         top_k = st.slider("Top K", 1, 30, 10, key="q_topk")
-    
-    prefetch_k, stage1_mode, stage1_k, stage2_k = 256, "tokens_vs_tiles", 1000, 300
-    
+
+    prefetch_k, stage1_mode, stage1_k, stage2_k = 256, "tokens_vs_standard_pooling", 1000, 300
+
     if mode == "two_stage":
         cc1, cc2 = st.columns(2)
         with cc1:
@@ -307,33 +316,44 @@ def render_rag_query_interface(collection_name: str, model_name: str = None):
             stage1_k = st.number_input("Stage1 K", 100, 5000, 1000, key="q_s1k")
         with cc2:
             stage2_k = st.number_input("Stage2 K", 50, 1000, 300, key="q_s2k")
-    
+
     with c1:
         query = st.text_input("Query", placeholder="Enter your search query...", key="q_text")
-    
+
     if st.button("🔍 Search", type="primary", disabled=not query, key="q_search"):
         with st.spinner("Searching..."):
             results, err = search_collection(
-                collection_name, query, top_k, mode, prefetch_k, stage1_mode, stage1_k, stage2_k, model_name
+                collection_name,
+                query,
+                top_k,
+                mode,
+                prefetch_k,
+                stage1_mode,
+                stage1_k,
+                stage2_k,
+                model_name,
             )
             if err:
                 st.error("Search failed")
                 st.code(err)
             else:
                 st.session_state["q_results"] = results
-    
+
     results = st.session_state.get("q_results", [])
     if results:
         st.success(f"**{len(results)} results**")
         max_score = max(r.get("score_final", r.get("score_stage1", 0)) for r in results) or 1
-        
+
         for i, r in enumerate(results):
             p = r.get("payload", {})
             score = r.get("score_final", r.get("score_stage1", 0))
             rel = score / max_score * 100
-            
+
             filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
             page_num = p.get("page_number") or p.get("page") or "?"
-            
-            with st.expander(f"**#{i+1}** {str(filename)[:35]} - Page {page_num} | 🎯 {rel:.0f}%", expanded=(i < 3)):
+
+            with st.expander(
+                f"**#{i+1}** {str(filename)[:35]} - Page {page_num} | 🎯 {rel:.0f}%",
+                expanded=(i < 3),
+            ):
                 render_document_details(r, p, score=score, rel_pct=rel)

demo/ui/sidebar.py

@@ -1,23 +1,24 @@
 """Sidebar component."""
 
 import os
-import streamlit as st
 
+import streamlit as st
 from qdrant_client.models import VectorParamsDiff
 
 from demo.qdrant_utils import (
+    get_collection_stats,
+    get_collections,
     get_qdrant_credentials,
+    get_vector_sizes,
     init_qdrant_client_with_creds,
-    get_collections,
-    get_collection_stats,
     sample_points_cached,
-    get_vector_sizes,
 )
 
 
 def render_sidebar():
     # CSS to make sidebar metrics smaller
-    st.markdown("""
+    st.markdown(
+        """
     <style>
     /* Smaller metrics in sidebar */
     [data-testid="stSidebar"] [data-testid="stMetricValue"] {
@@ -36,19 +37,21 @@ def render_sidebar():
         margin-bottom: 0.5rem !important;
     }
     </style>
-    """, unsafe_allow_html=True)
-    
+    """,
+        unsafe_allow_html=True,
+    )
+
     with st.sidebar:
         st.subheader("🔑 Qdrant Credentials")
-        
+
         env_url = os.getenv("QDRANT_URL") or os.getenv("SIGIR_QDRANT_URL") or ""
         env_key = os.getenv("QDRANT_API_KEY") or os.getenv("SIGIR_QDRANT_KEY") or ""
-        
+
         if "qdrant_url_input" not in st.session_state:
             st.session_state["qdrant_url_input"] = env_url
         if "qdrant_key_input" not in st.session_state:
             st.session_state["qdrant_key_input"] = env_key
-        
+
         qdrant_url = st.text_input(
             "Qdrant URL",
             value=st.session_state["qdrant_url_input"],
@@ -61,20 +64,23 @@ def render_sidebar():
             key="qdrant_key_widget",
             type="password",
         )
-        
-        if qdrant_url != st.session_state["qdrant_url_input"] or qdrant_key != st.session_state["qdrant_key_input"]:
+
+        if (
+            qdrant_url != st.session_state["qdrant_url_input"]
+            or qdrant_key != st.session_state["qdrant_key_input"]
+        ):
             st.session_state["qdrant_url_input"] = qdrant_url
             st.session_state["qdrant_key_input"] = qdrant_key
             get_collections.clear()
             get_collection_stats.clear()
             sample_points_cached.clear()
-        
+
         st.divider()
-        
+
         st.subheader("📡 Status")
         url, api_key = get_qdrant_credentials()
         client, err = init_qdrant_client_with_creds(url, api_key)
-        
+
         col_s1, col_s2 = st.columns(2)
         with col_s1:
             if client:
@@ -82,14 +88,16 @@ def render_sidebar():
             else:
                 st.error("Qdrant ✗", icon="❌")
         with col_s2:
-            cloudinary_ok = all([os.getenv("CLOUDINARY_CLOUD_NAME"), os.getenv("CLOUDINARY_API_KEY")])
+            cloudinary_ok = all(
+                [os.getenv("CLOUDINARY_CLOUD_NAME"), os.getenv("CLOUDINARY_API_KEY")]
+            )
             if cloudinary_ok:
                 st.success("Cloudinary ✓", icon="✅")
             else:
                 st.warning("Cloudinary ✗", icon="⚠️")
-        
+
         st.divider()
-        
+
         with st.expander("📦 Collection", expanded=True):
             collections = get_collections(url, api_key)
             if collections:
@@ -109,17 +117,23 @@ def render_sidebar():
                     if "error" not in stats:
                         col1, col2 = st.columns(2)
                         col1.metric("Points", f"{stats.get('points_count', 0):,}")
-                        status_raw = stats.get("status", "unknown").replace("CollectionStatus.", "").lower()
-                        status_icon = "🟢" if status_raw == "green" else "🟡" if status_raw == "yellow" else "🔴"
+                        status_raw = (
+                            stats.get("status", "unknown").replace("CollectionStatus.", "").lower()
+                        )
+                        status_icon = (
+                            "🟢"
+                            if status_raw == "green"
+                            else "🟡" if status_raw == "yellow" else "🔴"
+                        )
                         col2.metric("Status", status_icon)
-                        
+
                         points = stats.get("points_count", 0)
                         indexed = stats.get("indexed_vectors_count", 0) or 0
                         is_indexed = indexed >= points and points > 0
                         col3, col4 = st.columns(2)
                         col3.metric("Indexed", f"{indexed:,}")
                         col4.metric("HNSW", "✅" if is_indexed else "⏳")
-                        
+
                         vector_info = stats.get("vector_info", {})
                         if vector_info:
                             st.markdown("---")
@@ -135,14 +149,16 @@ def render_sidebar():
                                 on_disk = vinfo.get("on_disk", False)
                                 disk_icon = "💾" if on_disk else "🧠"
                                 dim_str = f"{num_vec}×{dim}"
-                                rows.append(f"<tr><td style='text-align:left;padding-right:12px;'><code>{vname}</code></td><td style='text-align:right;'>{dim_str}, {dtype}, {disk_icon}</td></tr>")
+                                rows.append(
+                                    f"<tr><td style='text-align:left;padding-right:12px;'><code>{vname}</code></td><td style='text-align:right;'>{dim_str}, {dtype}, {disk_icon}</td></tr>"
+                                )
                             table_html = f"<table style='width:100%;font-size:0.85em;'>{''.join(rows)}</table>"
                             st.markdown(table_html, unsafe_allow_html=True)
                     else:
                         st.error("Error loading stats")
             else:
                 st.info("No collections")
-        
+
         with st.expander("⚙️ Admin", expanded=False):
             active = st.session_state.get("active_collection")
             if active and client:
@@ -156,13 +172,17 @@ def render_sidebar():
                         current_on_disk = vector_info.get(sel_vec, {}).get("on_disk", False)
                         current_in_ram = not current_on_disk
                         st.caption(f"Current: {'🧠 RAM' if current_in_ram else '💾 Disk'}")
-                        target_in_ram = st.toggle("Move to RAM", value=current_in_ram, key=f"admin_ram_{sel_vec}")
+                        target_in_ram = st.toggle(
+                            "Move to RAM", value=current_in_ram, key=f"admin_ram_{sel_vec}"
+                        )
                         if target_in_ram != current_in_ram:
                             if st.button("💾 Apply Change", key="admin_apply"):
                                 try:
                                     client.update_collection(
                                         collection_name=active,
-                                        vectors_config={sel_vec: VectorParamsDiff(on_disk=not target_in_ram)}
+                                        vectors_config={
+                                            sel_vec: VectorParamsDiff(on_disk=not target_in_ram)
+                                        },
                                     )
                                     get_collection_stats.clear()
                                     st.success(f"Updated {sel_vec}")
@@ -175,9 +195,9 @@ def render_sidebar():
                     st.info("No vectors")
             else:
                 st.info("Select a collection")
-        
+
         st.divider()
-        
+
         if st.button("🔄 Refresh", type="secondary", use_container_width=True):
             get_collections.clear()
             get_collection_stats.clear()

demo/ui/upload.py

@@ -1,12 +1,11 @@
 """Upload tab component."""
 
+import inspect
+import json
 import os
 import tempfile
 import time
 import traceback
-import json
-import inspect
-from datetime import datetime
 from pathlib import Path
 
 import numpy as np
@@ -14,33 +13,33 @@ import streamlit as st
 
 from demo.config import AVAILABLE_MODELS
 from demo.qdrant_utils import (
-    get_qdrant_credentials,
     get_collection_stats,
+    get_qdrant_credentials,
     sample_points_cached,
 )
 from visual_rag.embedding.visual_embedder import VisualEmbedder
-from visual_rag.indexing.qdrant_indexer import QdrantIndexer
 from visual_rag.indexing.cloudinary_uploader import CloudinaryUploader
 from visual_rag.indexing.pipeline import ProcessingPipeline
-
+from visual_rag.indexing.qdrant_indexer import QdrantIndexer
 
 VECTOR_TYPES = ["initial", "mean_pooling", "experimental_pooling", "global_pooling"]
 
+
 def _load_metadata_mapping_from_uploaded_json(uploaded_json_file) -> tuple[dict, str]:
     """
     Load a filename->metadata mapping from an uploaded JSON file.
-    
+
     Supported formats:
     - Flat dict:
         { "Some Report 2023": {"year": 2023, "source": "...", ...}, ... }
     - Nested dict:
         { "filenames": { "Some Report 2023": {...}, ... }, ... }
-    
+
     Keys are normalized to: lowercase, trimmed, without ".pdf".
     """
     if uploaded_json_file is None:
         return {}, ""
-    
+
     try:
         raw = uploaded_json_file.getvalue()
         if not raw:
@@ -48,12 +47,12 @@ def _load_metadata_mapping_from_uploaded_json(uploaded_json_file) -> tuple[dict,
         data = json.loads(raw.decode("utf-8"))
         if not isinstance(data, dict):
             return {}, "Metadata file must be a JSON object"
-        
+
         mapping = data.get("filenames") if isinstance(data.get("filenames"), dict) else data
-        
+
         # Drop non-mapping keys (common pattern: _description, _usage)
         mapping = {k: v for k, v in mapping.items() if isinstance(k, str) and not k.startswith("_")}
-        
+
         normalized: dict[str, dict] = {}
         bad = 0
         for k, v in mapping.items():
@@ -67,7 +66,7 @@ def _load_metadata_mapping_from_uploaded_json(uploaded_json_file) -> tuple[dict,
                 bad += 1
                 continue
             normalized[key] = v
-        
+
         msg = f"Loaded {len(normalized):,} filename metadata mappings"
         if bad:
             msg += f" (ignored {bad:,} non-mapping entries)"
@@ -81,26 +80,30 @@ def render_upload_tab():
         msg = st.session_state.pop("upload_success")
         st.toast(f"✅ {msg}", icon="🎉")
         st.balloons()
-    
+
     st.subheader("📤 PDF Upload & Processing")
-    
+
     col_upload, col_config = st.columns([3, 2])
-    
+
     with col_config:
         st.markdown("##### Configuration")
-        
+
         c1, c2 = st.columns(2)
         with c1:
             model_name = st.selectbox("Model", AVAILABLE_MODELS, index=1, key="upload_model")
         with c2:
-            collection_name = st.text_input("Collection", value="my_collection", key="upload_collection_input")
-        
+            collection_name = st.text_input(
+                "Collection", value="my_collection", key="upload_collection_input"
+            )
+
         c3, c4 = st.columns(2)
         with c3:
-            vector_dtype = st.selectbox("Vector Dtype", ["float16", "float32"], index=0, key="upload_dtype")
+            vector_dtype = st.selectbox(
+                "Vector Dtype", ["float16", "float32"], index=0, key="upload_dtype"
+            )
         with c4:
             use_cloudinary = st.toggle("Cloudinary", value=True, key="upload_cloudinary")
-        
+
         st.markdown("**Performance**")
         p1, p2, p3 = st.columns(3)
         with p1:
@@ -139,9 +142,9 @@ def render_upload_tab():
             VECTOR_TYPES,
             default=VECTOR_TYPES,
             key="upload_vectors",
-            help="Which vector types to store in Qdrant"
+            help="Which vector types to store in Qdrant",
         )
-        
+
         st.markdown("**Crop Settings**")
         cc1, cc2 = st.columns(2)
         with cc1:
@@ -154,7 +157,9 @@ def render_upload_tab():
                 uniform_rowcol_std_threshold = st.select_slider(
                     "Uniform row/col threshold (any color)",
                     options=[0.0, 1.0, 2.0, 3.0, 5.0, 8.0, 12.0, 16.0],
-                    value=float(st.session_state.get("upload_uniform_rowcol_std_threshold", 0.0) or 0.0),
+                    value=float(
+                        st.session_state.get("upload_uniform_rowcol_std_threshold", 0.0) or 0.0
+                    ),
                     key="upload_uniform_rowcol_std_threshold",
                     help=(
                         "0 = off (default). Higher values remove more uniform borders, even if they are gray/black. "
@@ -165,13 +170,20 @@ def render_upload_tab():
                         "- 8+: aggressive (may remove faint content)"
                     ),
                 )
-        
+
         if crop_empty:
-            crop_pct = st.slider("Crop %", 0.90, 0.99, 0.99, 0.01, key="upload_crop_pct", 
-                                help="Remove margins with this % empty space")
+            crop_pct = st.slider(
+                "Crop %",
+                0.90,
+                0.99,
+                0.99,
+                0.01,
+                key="upload_crop_pct",
+                help="Remove margins with this % empty space",
+            )
         else:
             crop_pct = 0.99
-        
+
         st.markdown("**File Metadata (optional)**")
         meta_file = st.file_uploader(
             "Metadata mapping (JSON)",
@@ -195,7 +207,7 @@ def render_upload_tab():
                 st.warning(meta_msg or "No mappings loaded")
         else:
             metadata_mapping = {}
-    
+
     with col_upload:
         uploaded_files = st.file_uploader(
             "Select PDF files",
@@ -203,10 +215,10 @@ def render_upload_tab():
             accept_multiple_files=True,
             key="pdf_uploader",
         )
-        
+
         if uploaded_files:
             st.success(f"**{len(uploaded_files)} file(s) selected**")
-            
+
             if st.button("🚀 Process PDFs", type="primary", key="process_btn"):
                 config = {
                     "model_name": model_name,
@@ -223,7 +235,7 @@ def render_upload_tab():
                     "upload_batch_size": int(upload_batch_size),
                 }
                 process_pdfs(uploaded_files, config)
-    
+
     if st.session_state.get("last_upload_result"):
         st.divider()
         render_upload_results()
@@ -241,21 +253,21 @@ def process_pdfs(uploaded_files, config):
     dpi = int(config.get("dpi") or 140)
     embed_batch_size = int(config.get("embed_batch_size") or 8)
     upload_batch_size = int(config.get("upload_batch_size") or 8)
-    
+
     st.divider()
-    
+
     phase1 = st.container()
     phase2 = st.container()
     phase3 = st.container()
     results_container = st.container()
-    
+
     try:
         with phase1:
             st.markdown("##### 🤖 Phase 1: Loading Model")
             model_status = st.empty()
             model_short = model_name.split("/")[-1]
             model_status.info(f"Loading `{model_short}`...")
-            
+
             output_dtype = np.float16 if vector_dtype == "float16" else np.float32
             embedder_key = f"{model_name}::{vector_dtype}"
             embedder = None
@@ -267,18 +279,18 @@ def process_pdfs(uploaded_files, config):
                 st.session_state["upload_embedder_key"] = embedder_key
                 st.session_state["upload_embedder"] = embedder
             model_status.success(f"✅ Model `{model_short}` loaded ({vector_dtype})")
-        
+
         with phase2:
             st.markdown("##### 📦 Phase 2: Setting Up Collection")
-            
+
             url, api_key = get_qdrant_credentials()
             if not url or not api_key:
                 st.error("Qdrant credentials not configured")
                 return
-            
+
             qdrant_status = st.empty()
-            qdrant_status.info(f"Connecting to Qdrant...")
-            
+            qdrant_status.info("Connecting to Qdrant...")
+
             indexer = QdrantIndexer(
                 url=url,
                 api_key=api_key,
@@ -287,15 +299,15 @@ def process_pdfs(uploaded_files, config):
                 vector_datatype=vector_dtype,
                 timeout=180,
             )
-            qdrant_status.success(f"✅ Connected to Qdrant")
-            
+            qdrant_status.success("✅ Connected to Qdrant")
+
             coll_status = st.empty()
             collection_exists = False
             try:
                 collection_exists = indexer.collection_exists()
             except Exception:
                 pass
-            
+
             if collection_exists:
                 coll_status.success(f"✅ Collection `{collection_name}` exists (will append)")
             else:
@@ -304,24 +316,26 @@ def process_pdfs(uploaded_files, config):
                     try:
                         indexer.create_collection(force_recreate=False)
                         break
-                    except Exception as e:
+                    except Exception:
                         if attempt < 2:
                             time.sleep(2)
                         else:
                             raise
                 coll_status.success(f"✅ Collection `{collection_name}` created")
-            
+
             idx_status = st.empty()
             idx_status.info("Setting up indexes...")
             try:
-                indexer.create_payload_indexes(fields=[
-                    {"field": "filename", "type": "keyword"},
-                    {"field": "page_number", "type": "integer"},
-                ])
+                indexer.create_payload_indexes(
+                    fields=[
+                        {"field": "filename", "type": "keyword"},
+                        {"field": "page_number", "type": "integer"},
+                    ]
+                )
             except Exception:
                 pass
             idx_status.success("✅ Indexes ready")
-            
+
             cloud_status = st.empty()
             cloudinary_uploader = None
             if use_cloudinary:
@@ -333,9 +347,11 @@ def process_pdfs(uploaded_files, config):
                     cloud_status.warning(f"⚠️ Cloudinary unavailable: {str(e)[:30]}")
             else:
                 cloud_status.info("☁️ Cloudinary disabled")
-            
+
             pipeline = ProcessingPipeline(
-                embedder=embedder, indexer=indexer, cloudinary_uploader=cloudinary_uploader,
+                embedder=embedder,
+                indexer=indexer,
+                cloudinary_uploader=cloudinary_uploader,
                 metadata_mapping=metadata_mapping,
                 config={
                     "processing": {"dpi": dpi},
@@ -344,46 +360,54 @@ def process_pdfs(uploaded_files, config):
                         "upload_batch_size": upload_batch_size,
                     },
                 },
-                crop_empty=crop_empty, crop_empty_percentage_to_remove=crop_pct,
-                **({
-                    "crop_empty_uniform_rowcol_std_threshold": uniform_rowcol_std_threshold
-                } if "crop_empty_uniform_rowcol_std_threshold" in inspect.signature(ProcessingPipeline.__init__).parameters else {}),
+                crop_empty=crop_empty,
+                crop_empty_percentage_to_remove=crop_pct,
+                **(
+                    {"crop_empty_uniform_rowcol_std_threshold": uniform_rowcol_std_threshold}
+                    if "crop_empty_uniform_rowcol_std_threshold"
+                    in inspect.signature(ProcessingPipeline.__init__).parameters
+                    else {}
+                ),
             )
-        
+
         with phase3:
             st.markdown("##### 📄 Phase 3: Processing PDFs")
-            
+
             overall_progress = st.progress(0.0)
             file_status = st.empty()
             log_area = st.empty()
             log_lines = []
-            
+
             total_uploaded, total_skipped, total_failed = 0, 0, 0
             file_results = []
-            
+
             page_status = st.empty()
-            
+
             for i, f in enumerate(uploaded_files):
                 original_filename = f.name
-                file_status.info(f"📄 Processing `{original_filename}` ({i+1}/{len(uploaded_files)})")
+                file_status.info(
+                    f"📄 Processing `{original_filename}` ({i+1}/{len(uploaded_files)})"
+                )
                 t0 = time.perf_counter()
-                
+
                 with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
                     tmp.write(f.getvalue())
                     tmp_path = Path(tmp.name)
-                
+
                 def progress_cb(stage, current, total, message):
                     if stage == "process" and total > 0:
                         page_status.caption(f"  └─ Page {current}/{total}")
                     elif stage == "embed" and total > 0:
                         # Never show internal function names; keep this human-friendly.
-                        page_status.caption(f"  └─ Embedding pages… ({current+1}-{min(current + 1 + (pipeline.embedding_batch_size - 1), total)}/{total})")
+                        page_status.caption(
+                            f"  └─ Embedding pages… ({current+1}-{min(current + 1 + (pipeline.embedding_batch_size - 1), total)}/{total})"
+                        )
                     elif stage == "convert" and total > 0:
                         page_status.caption(f"  └─ {total} pages found")
-                
+
                 try:
                     result = pipeline.process_pdf(
-                        tmp_path, 
+                        tmp_path,
                         original_filename=original_filename,
                         progress_callback=progress_cb,
                     )
@@ -394,36 +418,46 @@ def process_pdfs(uploaded_files, config):
                     total_skipped += skipped
                     total_pages = int(result.get("total_pages") or 0)
                     sec_per_page = (elapsed_s / total_pages) if total_pages > 0 else None
-                    file_results.append({
-                        "file": original_filename,
-                        "uploaded": uploaded,
-                        "skipped": skipped,
-                        "total_pages": total_pages,
-                        "elapsed_s": float(elapsed_s),
-                        "sec_per_page": float(sec_per_page) if sec_per_page is not None else None,
-                    })
-                    timing_str = f"{elapsed_s:.1f}s" + (f" ({sec_per_page:.2f}s/page)" if sec_per_page is not None else "")
-                    log_lines.append(f"✓ {original_filename}: {uploaded} uploaded, {skipped} skipped | {timing_str}")
+                    file_results.append(
+                        {
+                            "file": original_filename,
+                            "uploaded": uploaded,
+                            "skipped": skipped,
+                            "total_pages": total_pages,
+                            "elapsed_s": float(elapsed_s),
+                            "sec_per_page": (
+                                float(sec_per_page) if sec_per_page is not None else None
+                            ),
+                        }
+                    )
+                    timing_str = f"{elapsed_s:.1f}s" + (
+                        f" ({sec_per_page:.2f}s/page)" if sec_per_page is not None else ""
+                    )
+                    log_lines.append(
+                        f"✓ {original_filename}: {uploaded} uploaded, {skipped} skipped | {timing_str}"
+                    )
                 except Exception as e:
                     total_failed += 1
                     log_lines.append(f"✗ {original_filename}: {str(e)[:50]}")
                 finally:
                     os.unlink(tmp_path)
-                
+
                 page_status.empty()
                 overall_progress.progress((i + 1) / len(uploaded_files))
                 log_area.code("\n".join(log_lines[-10:]), language="text")
-            
+
             overall_progress.progress(1.0)
             file_status.success(f"✅ Processed {len(uploaded_files)} file(s)")
-        
+
         with results_container:
             st.markdown("##### 📊 Results")
-            
-            st.success(f"✅ **{total_uploaded} pages** uploaded to `{collection_name}`" + 
-                       (f" ({total_skipped} skipped)" if total_skipped else "") +
-                       (f" ({total_failed} failed)" if total_failed else ""))
-            
+
+            st.success(
+                f"✅ **{total_uploaded} pages** uploaded to `{collection_name}`"
+                + (f" ({total_skipped} skipped)" if total_skipped else "")
+                + (f" ({total_failed} failed)" if total_failed else "")
+            )
+
             if file_results:
                 with st.expander("📋 File Details", expanded=False):
                     for fr in file_results:
@@ -438,19 +472,24 @@ def process_pdfs(uploaded_files, config):
                             + (f", {fr['total_pages']} pages" if fr.get("total_pages") else "")
                             + timing
                         )
-        
+
         st.session_state["last_upload_result"] = {
-            "total_uploaded": total_uploaded, "total_skipped": total_skipped, "total_failed": total_failed,
-            "file_results": file_results, "collection": collection_name,
+            "total_uploaded": total_uploaded,
+            "total_skipped": total_skipped,
+            "total_failed": total_failed,
+            "file_results": file_results,
+            "collection": collection_name,
         }
-        
+
         get_collection_stats.clear()
         sample_points_cached.clear()
-        
+
         if total_uploaded > 0:
-            st.session_state["upload_success"] = f"Uploaded {total_uploaded} pages to {collection_name}"
+            st.session_state["upload_success"] = (
+                f"Uploaded {total_uploaded} pages to {collection_name}"
+            )
             st.rerun()  # Immediately refresh to show success toast + balloons
-            
+
     except Exception as e:
         st.error(f"❌ Processing error: {e}")
         with st.expander("Traceback"):
@@ -461,17 +500,19 @@ def render_upload_results():
     result = st.session_state.get("last_upload_result", {})
     if not result:
         return
-    
+
     uploaded = result.get("total_uploaded", 0)
     skipped = result.get("total_skipped", 0)
     failed = result.get("total_failed", 0)
     collection = result.get("collection", "")
     file_results = result.get("file_results", [])
-    
-    st.success(f"✅ **{uploaded} pages** uploaded to `{collection}`" + 
-               (f" ({skipped} skipped)" if skipped else "") +
-               (f" ({failed} failed)" if failed else ""))
-    
+
+    st.success(
+        f"✅ **{uploaded} pages** uploaded to `{collection}`"
+        + (f" ({skipped} skipped)" if skipped else "")
+        + (f" ({failed} failed)" if failed else "")
+    )
+
     if file_results:
         with st.expander("📋 Details", expanded=False):
             for fr in file_results:

demo_app.py

@@ -0,0 +1,2068 @@
+import json
+import logging
+import os
+import tempfile
+import time
+import traceback
+import warnings
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+logging.getLogger("streamlit").setLevel(logging.ERROR)
+logging.getLogger("streamlit.runtime.scriptrunner_utils.script_run_context").setLevel(
+    logging.CRITICAL
+)
+warnings.filterwarnings("ignore", message=".*ScriptRunContext.*")
+
+os.environ.setdefault("STREAMLIT_SERVER_ENABLE_XSRF_PROTECTION", "false")
+os.environ.setdefault("STREAMLIT_SERVER_ENABLE_CORS", "false")
+os.environ.setdefault("STREAMLIT_SERVER_MAX_UPLOAD_SIZE", "500")
+os.environ.setdefault("STREAMLIT_BROWSER_GATHER_USAGE_STATS", "false")
+
+import altair as alt  # noqa: E402
+import numpy as np  # noqa: E402
+import pandas as pd  # noqa: E402
+import streamlit as st  # noqa: E402
+from dotenv import load_dotenv  # noqa: E402
+
+try:
+    from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+    from benchmarks.vidore_tatdqa_test.metrics import mrr_at_k, ndcg_at_k, recall_at_k
+    from visual_rag import VisualEmbedder
+    from visual_rag.indexing import QdrantIndexer
+    from visual_rag.retrieval import MultiVectorRetriever
+
+    VISUAL_RAG_AVAILABLE = True
+except ImportError:
+    VISUAL_RAG_AVAILABLE = False
+
+load_dotenv(Path(__file__).parent / ".env")
+if (Path(__file__).parent.parent / ".env").exists():
+    load_dotenv(Path(__file__).parent.parent / ".env")
+
+st.set_page_config(
+    page_title="Visual RAG Toolkit - Demo",
+    page_icon="🔬",
+    layout="wide",
+)
+
+AVAILABLE_MODELS = [
+    "vidore/colpali-v1.3",
+    "vidore/colSmol-500M",
+]
+
+BENCHMARK_DATASETS = [
+    "vidore/esg_reports_v2",
+    "vidore/biomedical_lectures_v2",
+    "vidore/economics_reports_v2",
+]
+
+DATASET_STATS = {
+    "vidore/esg_reports_v2": {"docs": 1538, "queries": 228},
+    "vidore/biomedical_lectures_v2": {"docs": 1016, "queries": 640},
+    "vidore/economics_reports_v2": {"docs": 452, "queries": 232},
+}
+
+RETRIEVAL_MODES = [
+    "single_full",
+    "single_tiles",
+    "single_global",
+    "two_stage",
+    "three_stage",
+]
+
+STAGE1_MODES = [
+    "tokens_vs_standard_pooling",
+    "tokens_vs_experimental_pooling",
+    "pooled_query_vs_standard_pooling",
+    "pooled_query_vs_experimental_pooling",
+    "pooled_query_vs_global",
+]
+
+
+def get_qdrant_credentials():
+    url = (
+        st.session_state.get("qdrant_url_input")
+        or os.getenv("SIGIR_QDRANT_URL")
+        or os.getenv("DEST_QDRANT_URL")
+        or os.getenv("QDRANT_URL")
+    )
+    api_key = st.session_state.get("qdrant_key_input") 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 url, api_key
+
+
+def init_qdrant_client_with_creds(url: str, api_key: str):
+    try:
+        from qdrant_client import QdrantClient
+
+        if not url:
+            return None, "QDRANT_URL not configured"
+        client = QdrantClient(url=url, api_key=api_key, timeout=60)
+        client.get_collections()
+        return client, None
+    except Exception as e:
+        return None, str(e)
+
+
+@st.cache_resource(show_spinner="Connecting to Qdrant...")
+def init_qdrant_client():
+    url, api_key = get_qdrant_credentials()
+    return init_qdrant_client_with_creds(url, api_key)
+
+
+@st.cache_resource(show_spinner="Loading embedding model...")
+def init_embedder(model_name: str):
+    try:
+        from visual_rag import VisualEmbedder
+
+        return VisualEmbedder(model_name=model_name), None
+    except Exception as e:
+        return None, f"{e}\n\n{traceback.format_exc()}"
+
+
+@st.cache_data(ttl=300, show_spinner="Fetching collections...")
+def get_collections(_url: str, _api_key: str) -> List[str]:
+    client, err = init_qdrant_client_with_creds(_url, _api_key)
+    if client is None:
+        return []
+    try:
+        collections = client.get_collections().collections
+        return sorted([c.name for c in collections])
+    except Exception:
+        return []
+
+
+@st.cache_data(ttl=120, show_spinner="Loading collection stats...")
+def get_collection_stats(collection_name: str) -> Dict[str, Any]:
+    url, api_key = get_qdrant_credentials()
+    client, err = init_qdrant_client_with_creds(url, api_key)
+    if client is None:
+        return {"error": err}
+    try:
+        info = client.get_collection(collection_name)
+        vectors_config = getattr(
+            getattr(getattr(info, "config", None), "params", None), "vectors", None
+        )
+        vector_info = {}
+        if vectors_config is not None:
+            if hasattr(vectors_config, "items"):
+                for name, cfg in vectors_config.items():
+                    size = getattr(cfg, "size", None)
+                    multivec = getattr(cfg, "multivector_config", None)
+                    on_disk = getattr(cfg, "on_disk", None)
+                    datatype = str(getattr(cfg, "datatype", "Float32")).replace("Datatype.", "")
+                    quantization = getattr(cfg, "quantization_config", None)
+                    num_vectors = 1
+                    if multivec is not None:
+                        comparator = getattr(multivec, "comparator", None)
+                        num_vectors = "N" if comparator else 1
+                    vector_info[name] = {
+                        "size": size,
+                        "num_vectors": num_vectors,
+                        "is_multivector": multivec is not None,
+                        "on_disk": on_disk,
+                        "datatype": datatype,
+                        "quantization": quantization is not None,
+                    }
+            elif hasattr(vectors_config, "size"):
+                on_disk = getattr(vectors_config, "on_disk", None)
+                datatype = str(getattr(vectors_config, "datatype", "Float32")).replace(
+                    "Datatype.", ""
+                )
+                multivec = getattr(vectors_config, "multivector_config", None)
+                vector_info["default"] = {
+                    "size": getattr(vectors_config, "size", None),
+                    "num_vectors": "N" if multivec else 1,
+                    "is_multivector": multivec is not None,
+                    "on_disk": on_disk,
+                    "datatype": datatype,
+                }
+        return {
+            "points_count": getattr(info, "points_count", 0),
+            "vectors_count": getattr(info, "vectors_count", getattr(info, "points_count", 0)),
+            "status": str(getattr(info, "status", "unknown")),
+            "vector_info": vector_info,
+            "indexed_vectors_count": getattr(info, "indexed_vectors_count", None),
+        }
+    except Exception as e:
+        return {"error": f"{e}\n\n{traceback.format_exc()}"}
+
+
+@st.cache_data(ttl=60)
+def sample_points_cached(
+    collection_name: str, n: int, seed: int, _url: str, _api_key: str
+) -> List[Dict[str, Any]]:
+    client, err = init_qdrant_client_with_creds(_url, _api_key)
+    if client is None:
+        return []
+    try:
+        import random
+
+        rng = random.Random(seed)
+        points, _ = client.scroll(
+            collection_name=collection_name,
+            limit=min(n * 10, 100),
+            with_payload=True,
+            with_vectors=False,
+        )
+        if not points:
+            return []
+        sampled = rng.sample(points, min(n, len(points)))
+        results = []
+        for p in sampled:
+            payload = dict(p.payload) if p.payload else {}
+            results.append(
+                {
+                    "id": str(p.id),
+                    "payload": payload,
+                }
+            )
+        return results
+    except Exception:
+        return []
+
+
+@st.cache_data(ttl=300)
+def get_vector_sizes(collection_name: str, _url: str, _api_key: str) -> Dict[str, int]:
+    client, err = init_qdrant_client_with_creds(_url, _api_key)
+    if client is None:
+        return {}
+    try:
+        points, _ = client.scroll(
+            collection_name=collection_name,
+            limit=1,
+            with_payload=False,
+            with_vectors=True,
+        )
+        if not points:
+            return {}
+        vectors = points[0].vector
+        sizes = {}
+        if isinstance(vectors, dict):
+            for name, vec in vectors.items():
+                if isinstance(vec, list):
+                    if vec and isinstance(vec[0], list):
+                        sizes[name] = len(vec)
+                    else:
+                        sizes[name] = 1
+                else:
+                    sizes[name] = 1
+        return sizes
+    except Exception:
+        return {}
+
+
+def search_collection(
+    collection_name: str,
+    query: str,
+    top_k: int = 10,
+    mode: str = "single_full",
+    prefetch_k: int = 256,
+    stage1_mode: str = "tokens_vs_standard_pooling",
+    stage1_k: int = 1000,
+    stage2_k: int = 300,
+    model_name: str = "vidore/colSmol-500M",
+) -> Tuple[List[Dict[str, Any]], Optional[str]]:
+    try:
+        from visual_rag.retrieval import MultiVectorRetriever
+
+        retriever = MultiVectorRetriever(
+            collection_name=collection_name,
+            model_name=model_name,
+        )
+        if mode == "three_stage":
+            q_emb = retriever.embedder.embed_query(query)
+            if hasattr(q_emb, "cpu"):
+                q_emb = q_emb.cpu().numpy()
+            results = retriever.search_embedded(
+                query_embedding=q_emb,
+                top_k=top_k,
+                mode=mode,
+                stage1_k=stage1_k,
+                stage2_k=stage2_k,
+            )
+        else:
+            results = retriever.search(
+                query=query,
+                top_k=top_k,
+                mode=mode,
+                prefetch_k=prefetch_k,
+                stage1_mode=stage1_mode,
+            )
+        return results, None
+    except Exception as e:
+        return [], f"{e}\n\n{traceback.format_exc()}"
+
+
+def load_results_file(path: Path) -> Optional[Dict[str, Any]]:
+    try:
+        with open(path, "r") as f:
+            return json.load(f)
+    except Exception:
+        return None
+
+
+def get_available_results() -> List[Path]:
+    results_dir = Path(__file__).parent / "results"
+    if not results_dir.exists():
+        return []
+    results = []
+    for subdir in results_dir.iterdir():
+        if subdir.is_dir():
+            for f in subdir.glob("*.json"):
+                if "index_failures" not in f.name:
+                    results.append(f)
+    return sorted(results, key=lambda x: x.stat().st_mtime, reverse=True)
+
+
+def find_main_result_file(collection: str, mode: str) -> Optional[Path]:
+    results = get_available_results()
+    for r in results:
+        if collection in str(r) and mode in r.name:
+            if "__vidore_" not in r.name:
+                return r
+    return results[0] if results else None
+
+
+def build_index_command(config: Dict[str, Any]) -> str:
+    cmd_parts = ["python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir"]
+    cmd_parts.append(f"--datasets {' '.join(config['datasets'])}")
+    cmd_parts.append(f"--collection {config['collection']}")
+    cmd_parts.append(f"--model {config['model']}")
+    cmd_parts.append("--index")
+    if config.get("recreate"):
+        cmd_parts.append("--recreate")
+    if config.get("resume"):
+        cmd_parts.append("--resume")
+    if config.get("prefer_grpc"):
+        cmd_parts.append("--prefer-grpc")
+    else:
+        cmd_parts.append("--no-prefer-grpc")
+    cmd_parts.append(f"--torch-dtype {config.get('torch_dtype', 'float16')}")
+    cmd_parts.append(f"--qdrant-vector-dtype {config.get('qdrant_vector_dtype', 'float16')}")
+    cmd_parts.append(f"--batch-size {config.get('batch_size', 4)}")
+    cmd_parts.append(f"--upload-batch-size {config.get('upload_batch_size', 8)}")
+    cmd_parts.append(f"--qdrant-timeout {config.get('qdrant_timeout', 180)}")
+    cmd_parts.append(f"--qdrant-retries {config.get('qdrant_retries', 5)}")
+    if config.get("crop_empty"):
+        cmd_parts.append("--crop-empty")
+        cmd_parts.append(f"--crop-empty-percentage-to-remove {config.get('crop_percentage', 0.99)}")
+    if config.get("no_cloudinary"):
+        cmd_parts.append("--no-cloudinary")
+    cmd_parts.append("--no-eval")
+    return " \\\n  ".join(cmd_parts)
+
+
+def build_eval_command(config: Dict[str, Any]) -> str:
+    cmd_parts = ["python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir"]
+    cmd_parts.append(f"--datasets {' '.join(config['datasets'])}")
+    cmd_parts.append(f"--collection {config['collection']}")
+    cmd_parts.append(f"--model {config['model']}")
+    cmd_parts.append(f"--mode {config['mode']}")
+    if config["mode"] == "two_stage":
+        cmd_parts.append(f"--stage1-mode {config.get('stage1_mode', 'tokens_vs_standard_pooling')}")
+        cmd_parts.append(f"--prefetch-k {config.get('prefetch_k', 256)}")
+    elif config["mode"] == "three_stage":
+        cmd_parts.append(f"--stage1-k {config.get('stage1_k', 1000)}")
+        cmd_parts.append(f"--stage2-k {config.get('stage2_k', 300)}")
+    cmd_parts.append(f"--top-k {config.get('top_k', 100)}")
+    cmd_parts.append(f"--evaluation-scope {config.get('evaluation_scope', 'union')}")
+    if config.get("prefer_grpc"):
+        cmd_parts.append("--prefer-grpc")
+    else:
+        cmd_parts.append("--no-prefer-grpc")
+    cmd_parts.append(f"--torch-dtype {config.get('torch_dtype', 'float16')}")
+    cmd_parts.append(f"--qdrant-vector-dtype {config.get('qdrant_vector_dtype', 'float16')}")
+    cmd_parts.append(f"--qdrant-timeout {config.get('qdrant_timeout', 180)}")
+    if config.get("result_prefix"):
+        cmd_parts.append(f"--output {config['result_prefix']}")
+    return " \\\n  ".join(cmd_parts)
+
+
+def generate_python_eval_code(config: Dict[str, Any]) -> str:
+    datasets_str = ", ".join([f'"{ds}"' for ds in config.get("datasets", [])])
+    mode = config.get("mode", "single_full")
+    model = config.get("model", "vidore/colpali-v1.3")
+    collection = config.get("collection", "")
+    top_k = config.get("top_k", 100)
+    scope = config.get("evaluation_scope", "union")
+    prefer_grpc = config.get("prefer_grpc", True)
+
+    code_lines = [
+        "import os",
+        "from qdrant_client import QdrantClient",
+        "from visual_rag import VisualEmbedder",
+        "from visual_rag.retrieval import MultiVectorRetriever",
+        "",
+        "# Configuration",
+        f'COLLECTION = "{collection}"',
+        f'MODEL = "{model}"',
+        f"TOP_K = {top_k}",
+        f"DATASETS = [{datasets_str}]",
+        "",
+        "# Initialize clients",
+        "client = QdrantClient(",
+        '    url=os.getenv("QDRANT_URL"),',
+        '    api_key=os.getenv("QDRANT_API_KEY"),',
+        f"    prefer_grpc={prefer_grpc},",
+        ")",
+        "",
+        "embedder = VisualEmbedder(",
+        "    model_name=MODEL,",
+        f'    torch_dtype="{config.get("torch_dtype", "float16")}",',
+        ")",
+        "",
+        "# Initialize retriever",
+        "retriever = MultiVectorRetriever(",
+        "    client=client,",
+        "    collection_name=COLLECTION,",
+        "    embedder=embedder,",
+        ")",
+        "",
+    ]
+
+    if mode == "single_full":
+        code_lines.extend(
+            [
+                "# Single-stage full retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="initial",',
+                "    )",
+            ]
+        )
+    elif mode == "single_tiles":
+        code_lines.extend(
+            [
+                "# Single-stage tiles retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="mean_pooling",',
+                "    )",
+            ]
+        )
+    elif mode == "single_global":
+        code_lines.extend(
+            [
+                "# Single-stage global retrieval",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return retriever.search_single_stage(",
+                "        query_embedding=query_embedding,",
+                f"        limit={top_k},",
+                '        vector_name="global_pooling",',
+                "    )",
+            ]
+        )
+    elif mode == "two_stage":
+        prefetch_k = config.get("prefetch_k", 256)
+        stage1_mode = config.get("stage1_mode", "tokens_vs_standard_pooling")
+        code_lines.extend(
+            [
+                "# Two-stage retrieval",
+                "from visual_rag.retrieval import TwoStageRetriever",
+                "",
+                "two_stage = TwoStageRetriever(",
+                "    client=client,",
+                "    collection_name=COLLECTION,",
+                "    embedder=embedder,",
+                ")",
+                "",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return two_stage.search(",
+                "        query_embedding=query_embedding,",
+                f"        prefetch_limit={prefetch_k},",
+                f"        limit={top_k},",
+                f'        stage1_mode="{stage1_mode}",',
+                "    )",
+            ]
+        )
+    elif mode == "three_stage":
+        stage1_k = config.get("stage1_k", 1000)
+        stage2_k = config.get("stage2_k", 300)
+        code_lines.extend(
+            [
+                "# Three-stage retrieval",
+                "from visual_rag.retrieval import ThreeStageRetriever",
+                "",
+                "three_stage = ThreeStageRetriever(",
+                "    client=client,",
+                "    collection_name=COLLECTION,",
+                "    embedder=embedder,",
+                ")",
+                "",
+                "def search(query: str):",
+                "    query_embedding = embedder.embed_query(query)",
+                "    return three_stage.search(",
+                "        query_embedding=query_embedding,",
+                f"        stage1_limit={stage1_k},",
+                f"        stage2_limit={stage2_k},",
+                f"        limit={top_k},",
+                "    )",
+            ]
+        )
+
+    if scope == "per_dataset":
+        code_lines.extend(
+            [
+                "",
+                "# Per-dataset filtering",
+                "from qdrant_client.models import Filter, FieldCondition, MatchValue",
+                "",
+                'def search_dataset(query: str, dataset: str = "vidore/esg_reports_v2"):',
+                "    query_embedding = embedder.embed_query(query)",
+                "    dataset_filter = Filter(",
+                "        must=[FieldCondition(",
+                '            key="dataset",',
+                "            match=MatchValue(value=dataset),",
+                "        )]",
+                "    )",
+                "    # Add filter to your search call",
+            ]
+        )
+
+    code_lines.extend(
+        [
+            "",
+            "# Example usage",
+            'results = search("What is the company revenue?")',
+            "for r in results:",
+            "    print(f\"Score: {r.score:.4f}, Doc: {r.payload.get('doc_id')}\")",
+        ]
+    )
+
+    return "\n".join(code_lines)
+
+
+def run_pythonic_evaluation(config: Dict[str, Any], progress_callback=None) -> Dict[str, Any]:
+    if not VISUAL_RAG_AVAILABLE:
+        raise ImportError("visual_rag package not available")
+
+    url, api_key = get_qdrant_credentials()
+    if not url:
+        raise ValueError("QDRANT_URL not configured")
+
+    datasets = config.get("datasets", [])
+    collection = config["collection"]
+    model = config.get("model", "vidore/colpali-v1.3")
+    mode = config.get("mode", "single_full")
+    top_k = config.get("top_k", 100)
+    prefetch_k = config.get("prefetch_k", 256)
+    stage1_mode = config.get("stage1_mode", "tokens_vs_standard_pooling")
+    stage1_k = config.get("stage1_k", 1000)
+    stage2_k = config.get("stage2_k", 300)
+    evaluation_scope = config.get("evaluation_scope", "union")
+    prefer_grpc = config.get("prefer_grpc", True)
+    torch_dtype = config.get("torch_dtype", "float16")
+
+    output_lines = []
+
+    def log(msg):
+        output_lines.append(msg)
+        if progress_callback:
+            progress_callback("\n".join(output_lines), None)
+
+    log(f"[Pythonic Eval] Initializing embedder: {model}")
+    embedder = VisualEmbedder(model_name=model, torch_dtype=torch_dtype)
+
+    log(f"[Pythonic Eval] Connecting to Qdrant collection: {collection}")
+    retriever = MultiVectorRetriever(
+        collection_name=collection,
+        model_name=model,
+        qdrant_url=url,
+        qdrant_api_key=api_key,
+        prefer_grpc=prefer_grpc,
+        embedder=embedder,
+    )
+
+    all_queries = []
+    all_qrels: Dict[str, Dict[str, int]] = {}
+    dataset_queries: Dict[str, List] = {}
+    dataset_qrels: Dict[str, Dict[str, Dict[str, int]]] = {}
+
+    for ds_name in datasets:
+        log(f"[Pythonic Eval] Loading dataset: {ds_name}")
+        corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
+        dataset_queries[ds_name] = queries
+        dataset_qrels[ds_name] = qrels
+        all_queries.extend(queries)
+        for qid, rels in qrels.items():
+            all_qrels[qid] = rels
+        log(f"  → {len(corpus)} docs, {len(queries)} queries")
+
+    def evaluate_queries(queries, qrels, filter_obj=None):
+        if not queries:
+            return {"ndcg@10": 0.0, "recall@10": 0.0, "mrr@10": 0.0, "num_queries": 0}
+
+        ndcg10_vals = []
+        recall10_vals = []
+        mrr10_vals = []
+        latencies = []
+
+        query_texts = [q.text for q in queries]
+        log(f"[Pythonic Eval] Embedding {len(query_texts)} queries...")
+        query_embeddings = embedder.embed_queries(query_texts, show_progress=False)
+
+        for i, (q, qemb) in enumerate(zip(queries, query_embeddings)):
+            start = time.time()
+
+            try:
+                import torch
+
+                if isinstance(qemb, torch.Tensor):
+                    qemb_np = qemb.detach().cpu().numpy()
+                else:
+                    qemb_np = qemb.numpy()
+            except ImportError:
+                qemb_np = qemb.numpy()
+
+            results = retriever.search_embedded(
+                query_embedding=qemb_np,
+                top_k=max(100, top_k),
+                mode=mode,
+                prefetch_k=prefetch_k,
+                stage1_mode=stage1_mode,
+                stage1_k=stage1_k,
+                stage2_k=stage2_k,
+                filter_obj=filter_obj,
+            )
+            latencies.append((time.time() - start) * 1000)
+
+            ranking = [str(r["id"]) for r in results]
+            rels = qrels.get(q.query_id, {})
+
+            ndcg10_vals.append(ndcg_at_k(ranking, rels, k=10))
+            recall10_vals.append(recall_at_k(ranking, rels, k=10))
+            mrr10_vals.append(mrr_at_k(ranking, rels, k=10))
+
+            if (i + 1) % 50 == 0:
+                log(f"  → Processed {i+1}/{len(queries)} queries")
+                if progress_callback:
+                    progress_callback("\n".join(output_lines), (i + 1) / len(queries))
+
+        return {
+            "ndcg@10": float(np.mean(ndcg10_vals)),
+            "recall@10": float(np.mean(recall10_vals)),
+            "mrr@10": float(np.mean(mrr10_vals)),
+            "avg_latency_ms": float(np.mean(latencies)),
+            "num_queries": len(queries),
+        }
+
+    results = {}
+
+    if evaluation_scope == "union":
+        log(f"\n[Pythonic Eval] Evaluating UNION ({len(all_queries)} queries)...")
+        union_metrics = evaluate_queries(all_queries, all_qrels)
+        results["union"] = union_metrics
+        log(f"  → NDCG@10: {union_metrics['ndcg@10']:.4f}")
+        log(f"  → Recall@10: {union_metrics['recall@10']:.4f}")
+        log(f"  → MRR@10: {union_metrics['mrr@10']:.4f}")
+    else:
+        for ds_name in datasets:
+            log(f"\n[Pythonic Eval] Evaluating {ds_name}...")
+            queries = dataset_queries[ds_name]
+            qrels = dataset_qrels[ds_name]
+            metrics = evaluate_queries(queries, qrels)
+            results[ds_name] = metrics
+            log(f"  → NDCG@10: {metrics['ndcg@10']:.4f}")
+            log(f"  → Recall@10: {metrics['recall@10']:.4f}")
+
+    log("\n" + "=" * 50)
+    log("[Pythonic Eval] COMPLETE!")
+
+    final_output = {
+        "config": {
+            "collection": collection,
+            "model": model,
+            "mode": mode,
+            "datasets": datasets,
+            "evaluation_scope": evaluation_scope,
+        },
+        "results": results,
+    }
+
+    return {"output": "\n".join(output_lines), "metrics": final_output}
+
+
+def run_pythonic_indexing(config: Dict[str, Any], progress_callback=None) -> Dict[str, Any]:
+    if not VISUAL_RAG_AVAILABLE:
+        raise ImportError("visual_rag package not available")
+
+    url, api_key = get_qdrant_credentials()
+    if not url:
+        raise ValueError("QDRANT_URL not configured")
+
+    datasets = config.get("datasets", [])
+    collection = config["collection"]
+    model = config.get("model", "vidore/colpali-v1.3")
+    recreate = config.get("recreate", False)
+    batch_size = config.get("batch_size", 4)
+    torch_dtype = config.get("torch_dtype", "float16")
+    qdrant_vector_dtype = config.get("qdrant_vector_dtype", "float16")
+    prefer_grpc = config.get("prefer_grpc", True)
+
+    output_lines = []
+
+    def log(msg):
+        output_lines.append(msg)
+        if progress_callback:
+            progress_callback("\n".join(output_lines), None)
+
+    log(f"[Pythonic Index] Initializing embedder: {model}")
+    embedder = VisualEmbedder(model_name=model, torch_dtype=torch_dtype)
+
+    log("[Pythonic Index] Connecting to Qdrant...")
+    indexer = QdrantIndexer(
+        url=url,
+        api_key=api_key,
+        collection_name=collection,
+        prefer_grpc=prefer_grpc,
+        vector_datatype=qdrant_vector_dtype,
+    )
+
+    log(f"[Pythonic Index] Creating collection: {collection}")
+    indexer.create_collection(force_recreate=recreate)
+
+    payload_fields = [
+        {"field": "dataset", "type": "keyword"},
+        {"field": "doc_id", "type": "keyword"},
+        {"field": "source_doc_id", "type": "keyword"},
+    ]
+    indexer.create_payload_indexes(fields=payload_fields)
+
+    total_uploaded = 0
+
+    for ds_name in datasets:
+        log(f"\n[Pythonic Index] Loading dataset: {ds_name}")
+        corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
+        log(f"  → {len(corpus)} documents to index")
+
+        for i in range(0, len(corpus), batch_size):
+            batch = corpus[i : i + batch_size]
+
+            images = []
+            for doc in batch:
+                img = doc.image if hasattr(doc, "image") else doc.get("image")
+                if img is not None:
+                    images.append(img)
+
+            if not images:
+                continue
+
+            log(
+                f"  → Embedding batch {i//batch_size + 1}/{(len(corpus) + batch_size - 1)//batch_size}..."
+            )
+            embeddings = embedder.embed_images(images)
+
+            points = []
+            for j, (doc, emb) in enumerate(zip(batch, embeddings)):
+                doc_id = doc.doc_id if hasattr(doc, "doc_id") else doc.get("doc_id", str(i + j))
+
+                if hasattr(emb, "cpu"):
+                    emb_np = emb.cpu().numpy()
+                else:
+                    emb_np = np.array(emb)
+
+                tile_pooled = emb_np.reshape(-1, 4, emb_np.shape[-1]).mean(axis=1)
+                global_pooled = emb_np.mean(axis=0)
+
+                points.append(
+                    {
+                        "id": f"{ds_name}_{doc_id}".replace("/", "_"),
+                        "visual_embedding": emb_np,
+                        "tile_pooled_embedding": tile_pooled,
+                        "experimental_pooled_embedding": tile_pooled,
+                        "global_pooled_embedding": global_pooled,
+                        "metadata": {
+                            "dataset": ds_name,
+                            "doc_id": doc_id,
+                            "source_doc_id": doc_id,
+                        },
+                    }
+                )
+
+            uploaded = indexer.upload_batch(points)
+            total_uploaded += uploaded
+
+            if progress_callback:
+                prog = (i + len(batch)) / len(corpus)
+                progress_callback("\n".join(output_lines), prog)
+
+        log(f"  → Finished {ds_name}: {total_uploaded} points uploaded")
+
+    log("\n" + "=" * 50)
+    log(f"[Pythonic Index] COMPLETE! Total: {total_uploaded} points")
+
+    return {"output": "\n".join(output_lines), "total_uploaded": total_uploaded}
+
+
+def render_header():
+    st.markdown(
+        """
+    <div style="text-align: center; padding: 10px 0 15px 0;">
+        <h1 style="
+            font-family: 'Georgia', serif;
+            font-size: 2.5rem;
+            font-weight: 700;
+            color: #1a1a2e;
+            letter-spacing: 3px;
+            margin: 0;
+            text-shadow: 2px 2px 4px rgba(0,0,0,0.1);
+        ">
+            🔬 Visual RAG Toolkit
+        </h1>
+        <p style="
+            font-family: 'Helvetica Neue', sans-serif;
+            font-size: 0.95rem;
+            color: #666;
+            margin-top: 5px;
+            letter-spacing: 1px;
+        ">
+            SIGIR 2026 Demo - Multi-Vector Visual Document Retrieval
+        </p>
+    </div>
+    """,
+        unsafe_allow_html=True,
+    )
+
+
+def render_sidebar():
+    with st.sidebar:
+        st.subheader("🔑 Qdrant Credentials")
+
+        env_url = (
+            os.getenv("SIGIR_QDRANT_URL")
+            or os.getenv("DEST_QDRANT_URL")
+            or os.getenv("QDRANT_URL")
+            or ""
+        )
+        env_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")
+            or ""
+        )
+
+        qdrant_url = st.text_input(
+            "Qdrant URL",
+            value=st.session_state.get("qdrant_url_input", env_url),
+            key="qdrant_url_widget",
+            placeholder="https://xxx.cloud.qdrant.io:6333",
+        )
+        qdrant_key = st.text_input(
+            "API Key",
+            value=st.session_state.get("qdrant_key_input", env_key),
+            key="qdrant_key_widget",
+            type="password",
+        )
+
+        if qdrant_url != st.session_state.get(
+            "qdrant_url_input"
+        ) or qdrant_key != st.session_state.get("qdrant_key_input"):
+            st.session_state["qdrant_url_input"] = qdrant_url
+            st.session_state["qdrant_key_input"] = qdrant_key
+            get_collections.clear()
+            get_collection_stats.clear()
+            sample_points_cached.clear()
+
+        st.divider()
+
+        st.subheader("📡 Status")
+        url, api_key = get_qdrant_credentials()
+        client, err = init_qdrant_client_with_creds(url, api_key)
+
+        col_s1, col_s2 = st.columns(2)
+        with col_s1:
+            if client:
+                st.success("Qdrant ✓", icon="✅")
+            else:
+                st.error("Qdrant ✗", icon="❌")
+        with col_s2:
+            cloudinary_ok = all(
+                [os.getenv("CLOUDINARY_CLOUD_NAME"), os.getenv("CLOUDINARY_API_KEY")]
+            )
+            if cloudinary_ok:
+                st.success("Cloudinary ✓", icon="✅")
+            else:
+                st.warning("Cloudinary ✗", icon="⚠️")
+
+        st.divider()
+
+        with st.expander("📦 Collection", expanded=True):
+            collections = get_collections(url, api_key)
+            if collections:
+                prev_collection = st.session_state.get("active_collection")
+                selected = st.selectbox(
+                    "Select Collection",
+                    options=collections,
+                    key="sidebar_collection",
+                    label_visibility="collapsed",
+                )
+                if selected:
+                    if selected != prev_collection:
+                        st.session_state["model_loaded"] = False
+                        st.session_state["loaded_model_key"] = None
+                    st.session_state["active_collection"] = selected
+                    stats = get_collection_stats(selected)
+                    if "error" not in stats:
+                        col1, col2 = st.columns(2)
+                        col1.metric("Points", f"{stats.get('points_count', 0):,}")
+                        status_raw = (
+                            stats.get("status", "unknown").replace("CollectionStatus.", "").lower()
+                        )
+                        status_icon = (
+                            "🟢"
+                            if status_raw == "green"
+                            else "🟡" if status_raw == "yellow" else "🔴"
+                        )
+                        col2.metric("Status", status_icon)
+
+                        points = stats.get("points_count", 0)
+                        indexed = stats.get("indexed_vectors_count", 0) or 0
+                        is_indexed = indexed >= points and points > 0
+                        col3, col4 = st.columns(2)
+                        col3.metric("Indexed", f"{indexed:,}")
+                        col4.metric("HNSW", "✅" if is_indexed else "⏳")
+
+                        vector_info = stats.get("vector_info", {})
+                        if vector_info:
+                            st.markdown("---")
+                            st.markdown("**🔢 Vectors**")
+                            vec_sizes = get_vector_sizes(selected, url, api_key)
+                            rows = []
+                            sorted_names = sorted(vector_info.keys(), key=lambda x: len(x))
+                            for vname in sorted_names:
+                                vinfo = vector_info[vname]
+                                dim = vinfo.get("size", "?")
+                                num_vec = vec_sizes.get(vname, vinfo.get("num_vectors", 1))
+                                dtype = vinfo.get("datatype", "?").upper()
+                                on_disk = vinfo.get("on_disk", False)
+                                disk_icon = "💾" if on_disk else "🧠"
+                                dim_str = f"{num_vec}×{dim}"
+                                rows.append(
+                                    f"<tr><td style='text-align:left;padding-right:12px;'><code>{vname}</code></td><td style='text-align:right;'>{dim_str}, {dtype}, {disk_icon}</td></tr>"
+                                )
+                            table_html = f"<table style='width:100%;font-size:0.85em;'>{''.join(rows)}</table>"
+                            st.markdown(table_html, unsafe_allow_html=True)
+                    else:
+                        st.error("Error loading stats")
+            else:
+                st.info("No collections")
+
+        with st.expander("⚙️ Admin", expanded=False):
+            active = st.session_state.get("active_collection")
+            if active and client:
+                stats = get_collection_stats(active)
+                vector_info = stats.get("vector_info", {})
+                if vector_info:
+                    st.markdown("**Change Storage**")
+                    vector_names = sorted(vector_info.keys())
+                    sel_vec = st.selectbox("Vector", vector_names, key="admin_vec")
+                    if sel_vec:
+                        current_on_disk = vector_info.get(sel_vec, {}).get("on_disk", False)
+                        current_in_ram = not current_on_disk
+                        st.caption(f"Current: {'🧠 RAM' if current_in_ram else '💾 Disk'}")
+                        target_in_ram = st.toggle(
+                            "Move to RAM", value=current_in_ram, key=f"admin_ram_{sel_vec}"
+                        )
+                        if target_in_ram != current_in_ram:
+                            if st.button("💾 Apply Change", key="admin_apply"):
+                                try:
+                                    from qdrant_client.models import VectorParamsDiff
+
+                                    client.update_collection(
+                                        collection_name=active,
+                                        vectors_config={
+                                            sel_vec: VectorParamsDiff(on_disk=not target_in_ram)
+                                        },
+                                    )
+                                    get_collection_stats.clear()
+                                    st.success(f"Updated {sel_vec}")
+                                    st.rerun()
+                                except Exception as e:
+                                    st.error(f"Failed: {e}")
+                        else:
+                            st.caption("Toggle to change storage location")
+                else:
+                    st.info("No vectors")
+            else:
+                st.info("Select a collection")
+
+        st.divider()
+
+        if st.button("🔄 Refresh", type="secondary", use_container_width=True):
+            get_collections.clear()
+            get_collection_stats.clear()
+            sample_points_cached.clear()
+            st.rerun()
+
+
+def render_upload_tab():
+    if "upload_success" in st.session_state:
+        msg = st.session_state.pop("upload_success")
+        st.toast(f"✅ {msg}", icon="🎉")
+        st.balloons()
+
+    st.subheader("📤 PDF Upload & Processing")
+
+    col_upload, col_config = st.columns([3, 2])
+
+    with col_config:
+        st.markdown("##### Configuration")
+
+        c1, c2 = st.columns(2)
+        with c1:
+            model_name = st.selectbox("Model", AVAILABLE_MODELS, index=1, key="upload_model")
+        with c2:
+            collection_name = st.text_input(
+                "Collection", value="my_collection", key="upload_collection_input"
+            )
+
+        c3, c4 = st.columns(2)
+        with c3:
+            crop_empty = st.toggle("Crop Margins", value=True, key="upload_crop")
+        with c4:
+            use_cloudinary = st.toggle("Cloudinary", value=True, key="upload_cloudinary")
+
+        if crop_empty:
+            crop_pct = st.slider("Crop %", 0.5, 0.99, 0.9, 0.01, key="upload_crop_pct")
+        else:
+            crop_pct = 0.9
+
+    with col_upload:
+        uploaded_files = st.file_uploader(
+            "Select PDF files",
+            type=["pdf"],
+            accept_multiple_files=True,
+            key="pdf_uploader",
+        )
+
+        if uploaded_files:
+            st.success(f"**{len(uploaded_files)} file(s) selected**")
+
+            if st.button("🚀 Process PDFs", type="primary", key="process_btn"):
+                process_pdfs(
+                    uploaded_files,
+                    model_name,
+                    collection_name,
+                    crop_empty,
+                    crop_pct,
+                    use_cloudinary,
+                )
+
+    if st.session_state.get("last_upload_result"):
+        st.divider()
+        render_upload_results()
+
+
+def process_pdfs(uploaded_files, model_name, collection_name, crop_empty, crop_pct, use_cloudinary):
+    logs = []
+    log_container = st.empty()
+    progress = st.progress(0)
+    status = st.empty()
+
+    def log(msg):
+        logs.append(f"[{datetime.now().strftime('%H:%M:%S')}] {msg}")
+        log_container.code("\n".join(logs[-30:]), language="text")
+
+    try:
+        log(f"Starting: {len(uploaded_files)} files, model={model_name.split('/')[-1]}")
+
+        from visual_rag import VisualEmbedder
+        from visual_rag.indexing import CloudinaryUploader, ProcessingPipeline, QdrantIndexer
+
+        log("Loading model...")
+        embedder = VisualEmbedder(model_name=model_name)
+
+        url, api_key = get_qdrant_credentials()
+        log("Connecting to Qdrant...")
+        indexer = QdrantIndexer(url=url, api_key=api_key, collection_name=collection_name)
+        indexer.create_collection(force_recreate=False)
+
+        cloudinary_uploader = None
+        if use_cloudinary:
+            try:
+                cloudinary_uploader = CloudinaryUploader()
+                log("Cloudinary ready")
+            except Exception as e:
+                log(f"Cloudinary failed: {e}")
+
+        pipeline = ProcessingPipeline(
+            embedder=embedder,
+            indexer=indexer,
+            cloudinary_uploader=cloudinary_uploader,
+            crop_empty=crop_empty,
+            crop_empty_percentage_to_remove=crop_pct,
+        )
+
+        total_uploaded, total_skipped, total_failed = 0, 0, 0
+        file_results = []
+
+        for i, f in enumerate(uploaded_files):
+            status.text(f"Processing: {f.name}")
+            log(f"[{i+1}/{len(uploaded_files)}] {f.name}")
+
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
+                tmp.write(f.getvalue())
+                tmp_path = Path(tmp.name)
+
+            try:
+                result = pipeline.process_pdf(tmp_path)
+                total_uploaded += result.get("uploaded", 0)
+                total_skipped += result.get("skipped", 0)
+                file_results.append(
+                    {
+                        "file": f.name,
+                        "uploaded": result.get("uploaded", 0),
+                        "skipped": result.get("skipped", 0),
+                    }
+                )
+                log(f"  ✓ uploaded={result.get('uploaded', 0)}, skipped={result.get('skipped', 0)}")
+            except Exception as e:
+                total_failed += 1
+                log(f"  ✗ Error: {e}")
+            finally:
+                os.unlink(tmp_path)
+
+            progress.progress((i + 1) / len(uploaded_files))
+
+        st.session_state["last_upload_result"] = {
+            "total_uploaded": total_uploaded,
+            "total_skipped": total_skipped,
+            "total_failed": total_failed,
+            "file_results": file_results,
+            "collection": collection_name,
+        }
+
+        get_collection_stats.clear()
+        sample_points_cached.clear()
+
+        if total_uploaded > 0:
+            st.session_state["upload_success"] = f"Uploaded {total_uploaded} pages"
+            st.rerun()
+
+    except Exception as e:
+        log(f"ERROR: {e}")
+        st.error(f"Processing error: {e}")
+        with st.expander("Traceback"):
+            st.code(traceback.format_exc())
+
+
+def render_upload_results():
+    result = st.session_state.get("last_upload_result", {})
+    if not result:
+        return
+
+    st.subheader("📊 Results")
+
+    c1, c2, c3 = st.columns(3)
+    c1.metric("Uploaded", result.get("total_uploaded", 0))
+    c2.metric("Skipped", result.get("total_skipped", 0))
+    c3.metric("Failed", result.get("total_failed", 0))
+
+
+def render_playground_tab():
+    st.subheader("🎮 Playground")
+
+    active_collection = st.session_state.get("active_collection")
+    url, api_key = get_qdrant_credentials()
+
+    if not active_collection:
+        collections = get_collections(url, api_key)
+        if collections:
+            active_collection = collections[0]
+
+    if not active_collection:
+        st.warning("No collection available. Upload documents or select a collection.")
+        return
+
+    points_for_model = sample_points_cached(active_collection, 1, 0, url, api_key)
+    model_name = None
+    if points_for_model:
+        model_name = points_for_model[0].get("payload", {}).get("model_name")
+    if not model_name:
+        model_name = AVAILABLE_MODELS[1]
+
+    model_short = model_name.split("/")[-1] if model_name else "unknown"
+    cache_key = f"{active_collection}_{model_name}"
+
+    if st.session_state.get("loaded_model_key") != cache_key:
+        st.session_state["model_loaded"] = False
+
+    col_info, col_model = st.columns([2, 1])
+    with col_info:
+        st.info(f"**Collection:** `{active_collection}`")
+    with col_model:
+        if not st.session_state.get("model_loaded"):
+            with st.spinner(f"Loading {model_short}..."):
+                try:
+                    from visual_rag.retrieval import MultiVectorRetriever
+
+                    _ = MultiVectorRetriever(
+                        collection_name=active_collection, model_name=model_name
+                    )
+                    st.session_state["model_loaded"] = True
+                    st.session_state["loaded_model_key"] = cache_key
+                    st.session_state["loaded_model_name"] = model_name
+                except Exception:
+                    st.warning(f"Failed: {model_short}")
+
+        if st.session_state.get("model_loaded"):
+            st.markdown(
+                f"✅ Found <span style='color:#e74c3c;font-weight:bold;'>{model_short}</span> model",
+                unsafe_allow_html=True,
+            )
+
+    with st.expander("📦 Sample Points Explorer", expanded=True):
+        render_sample_explorer(active_collection, url, api_key)
+
+    st.divider()
+
+    st.subheader("🔍 RAG Query")
+    render_rag_query_interface(active_collection, model_name)
+
+
+def render_document_details(pt: dict, p: dict, score: float = None, rel_pct: float = None):
+    doc_id = p.get("doc_id") or p.get("union_doc_id") or p.get("source_doc_id") or "?"
+    corpus_id = p.get("corpus-id") or p.get("source_doc_id") or "?"
+    dataset = p.get("dataset") or p.get("source") or "N/A"
+    model = (p.get("model_name") or p.get("model") or "N/A").split("/")[-1]
+    doc_name = p.get("doc-id") or p.get("filename") or "Unknown"
+
+    num_tiles = p.get("num_tiles") or "?"
+    visual_tokens = p.get("index_recovery_num_visual_tokens") or p.get("num_visual_tokens") or "?"
+    patches_per_tile = p.get("patches_per_tile") or "?"
+    torch_dtype = p.get("torch_dtype") or "?"
+
+    orig_w = p.get("original_width") or "?"
+    orig_h = p.get("original_height") or "?"
+    crop_w = p.get("cropped_width") or "?"
+    crop_h = p.get("cropped_height") or "?"
+    resize_w = p.get("resized_width") or "?"
+    resize_h = p.get("resized_height") or "?"
+    crop_pct = p.get("crop_empty_percentage_to_remove") or 0
+    crop_enabled = p.get("crop_empty_enabled", False)
+
+    col_meta, col_img = st.columns([1, 2])
+
+    with col_meta:
+        st.markdown("##### 📄 Document Info")
+        st.markdown(f"**📁 Doc:** {doc_name}")
+        st.markdown(f"**🏛️ Dataset:** {dataset}")
+        st.markdown(f"**🔑 Doc ID:** `{str(doc_id)[:20]}...`")
+        st.markdown(f"**📋 Corpus ID:** {corpus_id}")
+
+        if score is not None:
+            st.divider()
+            st.markdown("##### 🎯 Relevance")
+            if rel_pct is not None:
+                st.markdown(f"**Relative:** 🟢 {rel_pct:.1f}%")
+                st.progress(rel_pct / 100)
+            st.caption(f"Raw score: {score:.4f}")
+
+        st.divider()
+        st.markdown("##### 🎨 Visual Metadata")
+        st.markdown(f"**🤖 Model:** `{model}`")
+        st.markdown(f"**🔲 Tiles:** {num_tiles}")
+        st.markdown(f"**🔢 Visual Tokens:** {visual_tokens}")
+        st.markdown(f"**📦 Patches/Tile:** {patches_per_tile}")
+        st.markdown(f"**⚙️ Dtype:** {torch_dtype}")
+
+        st.divider()
+        st.markdown("##### 📐 Dimensions")
+        st.markdown(f"**Original:** {orig_w}×{orig_h}")
+        st.markdown(f"**Resized:** {resize_w}×{resize_h}")
+        if crop_enabled:
+            st.markdown(f"**Cropped:** {crop_w}×{crop_h}")
+            st.markdown(f"**Crop %:** {int(crop_pct * 100) if crop_pct else 0}%")
+
+    with col_img:
+        st.markdown("##### 📷 Document Page")
+        tabs = st.tabs(["🖼️ Original", "📷 Resized", "✂️ Cropped"])
+
+        url_o = p.get("original_url")
+        url_r = p.get("resized_url") or p.get("page")
+        url_c = p.get("cropped_url")
+
+        with tabs[0]:
+            if url_o:
+                st.image(url_o, width=600)
+                st.caption(f"📐 **{orig_w}×{orig_h}**")
+            else:
+                st.info("No original image available")
+
+        with tabs[1]:
+            if url_r:
+                st.image(url_r, width=600)
+                st.caption(f"📐 **{resize_w}×{resize_h}**")
+            else:
+                st.info("No resized image available")
+
+        with tabs[2]:
+            if url_c:
+                st.image(url_c, width=600)
+                st.caption(
+                    f"📐 **{crop_w}×{crop_h}** | Crop: {int(crop_pct * 100) if crop_pct else 0}%"
+                )
+            else:
+                st.info("No cropped image available")
+
+        with st.expander("🔗 Image URLs"):
+            if url_o:
+                st.code(url_o, language=None)
+            if url_r and url_r != url_o:
+                st.code(url_r, language=None)
+            if url_c:
+                st.code(url_c, language=None)
+
+
+def render_sample_explorer(collection_name: str, url: str, api_key: str):
+    sample_for_filters = sample_points_cached(collection_name, 50, 0, url, api_key)
+    datasets = set()
+    doc_ids = set()
+    for pt in sample_for_filters:
+        p = pt.get("payload", {})
+        if ds := p.get("dataset"):
+            datasets.add(ds)
+        if did := (p.get("doc-id") or p.get("filename")):
+            doc_ids.add(did)
+
+    c1, c2, c3, c4 = st.columns([1, 1, 2, 1])
+    with c1:
+        n_samples = st.slider("Samples", 1, 20, 3, key="pg_n")
+    with c2:
+        seed = st.number_input("Seed", 0, 9999, 42, key="pg_seed")
+    with c3:
+        filter_ds = st.selectbox("Dataset", ["All"] + sorted(datasets), key="pg_filter_ds")
+    with c4:
+        st.write("")
+        do_sample = st.button("🎲 Sample", type="primary", key="pg_sample_btn")
+
+    if do_sample:
+        points = sample_points_cached(collection_name, n_samples * 5, seed, url, api_key)
+        if filter_ds != "All":
+            points = [p for p in points if p.get("payload", {}).get("dataset") == filter_ds]
+        points = points[:n_samples]
+        st.session_state["pg_points"] = points
+
+    points = st.session_state.get("pg_points", [])
+
+    if not points:
+        st.caption("Click 'Sample' to load documents")
+        return
+
+    st.success(f"**{len(points)} points loaded**")
+
+    for i, pt in enumerate(points):
+        p = pt.get("payload", {})
+
+        filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
+        page_num = p.get("page_number") or p.get("page") or "?"
+
+        with st.expander(f"**{i+1}.** {str(filename)[:40]} - Page {page_num}", expanded=(i == 0)):
+            render_document_details(pt, p)
+
+
+def render_rag_query_interface(collection_name: str, model_name: str = None):
+    if not collection_name:
+        return
+
+    url, api_key = get_qdrant_credentials()
+
+    if not model_name:
+        points = sample_points_cached(collection_name, 1, 0, url, api_key)
+        if points:
+            model_name = points[0].get("payload", {}).get("model_name")
+        if not model_name:
+            model_name = AVAILABLE_MODELS[1]
+
+    st.caption(f"Model: **{model_name.split('/')[-1] if model_name else 'auto'}**")
+
+    c1, c2, c3 = st.columns([2, 1, 1])
+    with c2:
+        mode = st.selectbox("Mode", RETRIEVAL_MODES, index=0, key="q_mode")
+    with c3:
+        top_k = st.slider("Top K", 1, 30, 10, key="q_topk")
+
+    prefetch_k, stage1_mode, stage1_k, stage2_k = 256, "tokens_vs_standard_pooling", 1000, 300
+
+    if mode == "two_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_mode = st.selectbox("Stage1", STAGE1_MODES, key="q_s1mode")
+        with cc2:
+            prefetch_k = st.slider("Prefetch K", 50, 500, 256, key="q_pk")
+    elif mode == "three_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_k = st.number_input("Stage1 K", 100, 5000, 1000, key="q_s1k")
+        with cc2:
+            stage2_k = st.number_input("Stage2 K", 50, 1000, 300, key="q_s2k")
+
+    with c1:
+        query = st.text_input("Query", placeholder="Enter your search query...", key="q_text")
+
+    if st.button("🔍 Search", type="primary", disabled=not query, key="q_search"):
+        with st.spinner("Searching..."):
+            results, err = search_collection(
+                collection_name,
+                query,
+                top_k,
+                mode,
+                prefetch_k,
+                stage1_mode,
+                stage1_k,
+                stage2_k,
+                model_name,
+            )
+            if err:
+                st.error("Search failed")
+                st.code(err)
+            else:
+                st.session_state["q_results"] = results
+
+    results = st.session_state.get("q_results", [])
+    if results:
+        st.success(f"**{len(results)} results**")
+        max_score = max(r.get("score_final", r.get("score_stage1", 0)) for r in results) or 1
+
+        for i, r in enumerate(results):
+            p = r.get("payload", {})
+            score = r.get("score_final", r.get("score_stage1", 0))
+            rel = score / max_score * 100
+
+            filename = p.get("filename") or p.get("doc_id") or p.get("source_doc_id") or "Unknown"
+            page_num = p.get("page_number") or p.get("page") or "?"
+
+            with st.expander(
+                f"**#{i+1}** {str(filename)[:35]} - Page {page_num} | 🎯 {rel:.0f}%",
+                expanded=(i < 3),
+            ):
+                render_document_details(r, p, score=score, rel_pct=rel)
+
+
+def render_benchmark_tab():
+    st.subheader("📊 Benchmarking")
+
+    tab_index, tab_eval, tab_results = st.tabs(["Indexing", "Evaluation", "Results"])
+
+    url, api_key = get_qdrant_credentials()
+    collections = get_collections(url, api_key)
+
+    with tab_index:
+        render_benchmark_indexing(collections)
+
+    with tab_eval:
+        render_benchmark_evaluation(collections)
+
+    with tab_results:
+        render_benchmark_results()
+
+
+def render_benchmark_indexing(collections: List[str]):
+    c1, c2, c3 = st.columns(3)
+    with c1:
+        datasets = st.multiselect(
+            "Datasets", BENCHMARK_DATASETS, default=BENCHMARK_DATASETS, key="bi_ds"
+        )
+    with c2:
+        model = st.selectbox("Model", AVAILABLE_MODELS, key="bi_model")
+    with c3:
+        model_short = model.split("/")[-1].replace("-", "_").replace(".", "_")
+        collection = st.text_input(
+            "Collection", value=f"vidore_{len(datasets)}ds__{model_short}", key="bi_coll"
+        )
+
+    c4, c5, c6, c7 = st.columns(4)
+    with c4:
+        crop = st.toggle("Crop", value=True, key="bi_crop")
+    with c5:
+        cloudinary = st.toggle("Cloudinary", value=True, key="bi_cloud")
+    with c6:
+        grpc = st.toggle("gRPC", value=True, key="bi_grpc")
+    with c7:
+        recreate = st.toggle("Recreate", value=False, key="bi_recreate")
+
+    crop_pct = st.slider("Crop %", 0.8, 0.99, 0.99, 0.01, key="bi_crop_pct") if crop else 0.99
+
+    config = {
+        "datasets": datasets,
+        "model": model,
+        "collection": collection,
+        "crop_empty": crop,
+        "crop_percentage": crop_pct,
+        "no_cloudinary": not cloudinary,
+        "recreate": recreate,
+        "resume": False,
+        "prefer_grpc": grpc,
+        "batch_size": 4,
+        "upload_batch_size": 8,
+        "qdrant_timeout": 180,
+        "qdrant_retries": 5,
+        "torch_dtype": "float16",
+        "qdrant_vector_dtype": "float16",
+    }
+
+    cmd = build_index_command(config)
+
+    col_cmd, col_stats = st.columns([2, 1])
+    with col_cmd:
+        st.code(cmd, language="bash")
+    with col_stats:
+        st.metric("Datasets", len(datasets))
+        st.metric("Model", model.split("/")[-1])
+        run_index = st.button("🚀 Run Index", type="primary", key="bi_run")
+
+    if run_index:
+        if not collection:
+            st.error("Please select a collection first")
+        else:
+            run_indexing_with_ui(config)
+
+
+def render_benchmark_evaluation(collections: List[str]):
+    all_docs = sum(DATASET_STATS.get(d, {}).get("docs", 0) for d in BENCHMARK_DATASETS)
+    all_queries = sum(DATASET_STATS.get(d, {}).get("queries", 0) for d in BENCHMARK_DATASETS)
+    st.markdown(
+        f"📊 **Available:** {len(BENCHMARK_DATASETS)} datasets — **{all_docs:,}** docs, **{all_queries:,}** queries"
+    )
+
+    c1, c2, c3 = st.columns([2, 2, 1])
+    with c1:
+        if collections:
+            collection = st.selectbox("Collection", collections, key="be_coll")
+        else:
+            collection = st.text_input("Collection", key="be_coll_txt")
+    with c2:
+        st.multiselect("Datasets", BENCHMARK_DATASETS, default=BENCHMARK_DATASETS, key="be_ds")
+    with c3:
+        model = st.selectbox("Model", AVAILABLE_MODELS, key="be_model")
+
+    datasets = st.session_state.get("be_ds", BENCHMARK_DATASETS)
+    sel_docs = sum(DATASET_STATS.get(d, {}).get("docs", 0) for d in datasets)
+    sel_queries = sum(DATASET_STATS.get(d, {}).get("queries", 0) for d in datasets)
+    st.markdown(
+        f"🎯 **Selected:** {len(datasets)} dataset(s) — **{sel_docs:,}** docs, **{sel_queries:,}** queries"
+    )
+
+    st.markdown("---")
+
+    col_mode, col_topk = st.columns([2, 1])
+    with col_mode:
+        mode = st.selectbox("Mode", RETRIEVAL_MODES, key="be_mode")
+    with col_topk:
+        top_k = st.slider("Top K", 10, 100, 100, key="be_topk")
+
+    stage1_mode, prefetch_k, stage1_k, stage2_k = "tokens_vs_standard_pooling", 256, 1000, 300
+
+    if mode == "two_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_mode = st.selectbox("Stage1 Mode", STAGE1_MODES, key="be_s1mode")
+        with cc2:
+            prefetch_k = st.slider("Prefetch K", 50, 1000, 256, key="be_pk")
+    elif mode == "three_stage":
+        cc1, cc2 = st.columns(2)
+        with cc1:
+            stage1_k = st.number_input("Stage1 K", 100, 5000, 1000, key="be_s1k")
+        with cc2:
+            stage2_k = st.number_input("Stage2 K", 50, 1000, 300, key="be_s2k")
+
+    st.markdown("---")
+
+    col_scope, col_grpc, col_spacer = st.columns([2, 1, 1])
+    with col_scope:
+        scope = st.selectbox("Scope", ["union", "per_dataset"], key="be_scope")
+    with col_grpc:
+        grpc = st.toggle("gRPC", value=True, key="be_grpc")
+
+    result_prefix_val = st.session_state.get("be_prefix", "")
+
+    config = {
+        "datasets": datasets,
+        "model": model,
+        "collection": collection,
+        "mode": mode,
+        "top_k": top_k,
+        "evaluation_scope": scope,
+        "prefer_grpc": grpc,
+        "torch_dtype": "float16",
+        "qdrant_vector_dtype": "float16",
+        "qdrant_timeout": 180,
+        "stage1_mode": stage1_mode,
+        "prefetch_k": prefetch_k,
+        "stage1_k": stage1_k,
+        "stage2_k": stage2_k,
+        "result_prefix": result_prefix_val,
+    }
+
+    cmd = build_eval_command(config)
+
+    python_code = generate_python_eval_code(config)
+
+    col_cmd, col_info = st.columns([2, 1])
+    with col_cmd:
+        code_tab1, code_tab2 = st.tabs(["🐚 Bash", "🐍 Python"])
+        with code_tab1:
+            st.code(cmd, language="bash")
+        with code_tab2:
+            st.code(python_code, language="python")
+    with col_info:
+        mode_desc = {
+            "single_full": "🔹 **Single Full**: Query all visual tokens against full document embeddings in one pass.",
+            "single_tiles": "🔸 **Single Tiles**: Query against tile-level embeddings only.",
+            "single_global": "🔶 **Single Global**: Query against global (pooled) document embeddings.",
+            "two_stage": "🔷 **Two Stage**: Fast prefetch with global/tiles, then rerank with full tokens.",
+            "three_stage": "🔶 **Three Stage**: Global → Tiles → Full tokens for maximum precision.",
+        }
+        scope_desc = {
+            "union": "📊 **Union**: Evaluate across all datasets combined as one corpus.",
+            "per_dataset": "📁 **Per Dataset**: Evaluate each dataset separately and report individual metrics.",
+        }
+        st.markdown(mode_desc.get(mode, ""))
+        st.markdown(scope_desc.get(scope, ""))
+        st.divider()
+        st.text_input("Result Prefix", placeholder="optional prefix for output", key="be_prefix")
+
+        run_eval = st.button("🚀 Run Eval", type="primary", key="be_run", use_container_width=True)
+
+    if run_eval:
+        if not collection:
+            st.error("Please select a collection first")
+        else:
+            run_evaluation_with_ui(config)
+
+
+def run_evaluation_with_ui(config: Dict[str, Any]):
+    st.divider()
+
+    progress_bar = st.progress(0.0)
+    status_text = st.empty()
+    output_area = st.empty()
+
+    status_text.info("🚀 Starting evaluation...")
+    output_lines = []
+
+    def log(msg):
+        output_lines.append(msg)
+        output_area.code("\n".join(output_lines[-50:]), language="text")
+
+    try:
+        url, api_key = get_qdrant_credentials()
+        if not url:
+            st.error("QDRANT_URL not configured")
+            return
+
+        datasets = config.get("datasets", [])
+        collection = config["collection"]
+        model = config.get("model", "vidore/colpali-v1.3")
+        mode = config.get("mode", "single_full")
+        top_k = config.get("top_k", 100)
+        prefetch_k = config.get("prefetch_k", 256)
+        stage1_mode = config.get("stage1_mode", "tokens_vs_standard_pooling")
+        stage1_k = config.get("stage1_k", 1000)
+        stage2_k = config.get("stage2_k", 300)
+        _evaluation_scope = config.get("evaluation_scope", "union")
+        prefer_grpc = config.get("prefer_grpc", True)
+        torch_dtype = config.get("torch_dtype", "float16")
+
+        log(f"[Eval] Model: {model}")
+        log(f"[Eval] Collection: {collection}")
+        log(f"[Eval] Mode: {mode}")
+        log(f"[Eval] Datasets: {datasets}")
+        status_text.info("📦 Loading embedder...")
+
+        embedder = VisualEmbedder(model_name=model, torch_dtype=torch_dtype)
+        log("[Eval] Embedder loaded")
+
+        status_text.info("🔌 Connecting to Qdrant...")
+        retriever = MultiVectorRetriever(
+            collection_name=collection,
+            model_name=model,
+            qdrant_url=url,
+            qdrant_api_key=api_key,
+            prefer_grpc=prefer_grpc,
+            embedder=embedder,
+        )
+        log("[Eval] Retriever connected")
+
+        all_queries = []
+        all_qrels: Dict[str, Dict[str, int]] = {}
+
+        for ds_name in datasets:
+            status_text.info(f"📚 Loading dataset: {ds_name}")
+            corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
+            all_queries.extend(queries)
+            for qid, rels in qrels.items():
+                all_qrels[qid] = rels
+            log(f"[Eval] Loaded {ds_name}: {len(corpus)} docs, {len(queries)} queries")
+
+        total_queries = len(all_queries)
+        log(f"[Eval] Total queries to evaluate: {total_queries}")
+
+        status_text.info(f"🔍 Embedding {total_queries} queries...")
+        query_texts = [q.text for q in all_queries]
+        query_embeddings = embedder.embed_queries(query_texts, show_progress=False)
+        log("[Eval] Queries embedded")
+
+        ndcg10_vals = []
+        recall10_vals = []
+        mrr10_vals = []
+        latencies = []
+
+        status_text.info("🎯 Running evaluation...")
+
+        for i, (q, qemb) in enumerate(zip(all_queries, query_embeddings)):
+            start = time.time()
+
+            try:
+                import torch
+
+                if isinstance(qemb, torch.Tensor):
+                    qemb_np = qemb.detach().cpu().numpy()
+                else:
+                    qemb_np = qemb.numpy()
+            except ImportError:
+                qemb_np = qemb.numpy()
+
+            results = retriever.search_embedded(
+                query_embedding=qemb_np,
+                top_k=max(100, top_k),
+                mode=mode,
+                prefetch_k=prefetch_k,
+                stage1_mode=stage1_mode,
+                stage1_k=stage1_k,
+                stage2_k=stage2_k,
+            )
+            latencies.append((time.time() - start) * 1000)
+
+            ranking = [str(r["id"]) for r in results]
+            rels = all_qrels.get(q.query_id, {})
+
+            ndcg10_vals.append(ndcg_at_k(ranking, rels, k=10))
+            recall10_vals.append(recall_at_k(ranking, rels, k=10))
+            mrr10_vals.append(mrr_at_k(ranking, rels, k=10))
+
+            progress = (i + 1) / total_queries
+            progress_bar.progress(progress)
+            status_text.info(f"🎯 Evaluating... {i+1}/{total_queries} ({int(progress*100)}%)")
+
+            if (i + 1) % 20 == 0:
+                log(f"[Eval] Progress: {i+1}/{total_queries} queries")
+
+        progress_bar.progress(1.0)
+        status_text.success("✅ Evaluation complete!")
+
+        final_metrics = {
+            "ndcg@10": float(np.mean(ndcg10_vals)),
+            "recall@10": float(np.mean(recall10_vals)),
+            "mrr@10": float(np.mean(mrr10_vals)),
+            "avg_latency_ms": float(np.mean(latencies)),
+            "num_queries": total_queries,
+        }
+
+        log("")
+        log("=" * 40)
+        log("RESULTS:")
+        log(f"  NDCG@10:  {final_metrics['ndcg@10']:.4f}")
+        log(f"  Recall@10: {final_metrics['recall@10']:.4f}")
+        log(f"  MRR@10:   {final_metrics['mrr@10']:.4f}")
+        log(f"  Avg Latency: {final_metrics['avg_latency_ms']:.1f}ms")
+        log("=" * 40)
+
+        st.json(final_metrics)
+        st.session_state["last_eval_metrics"] = final_metrics
+
+    except Exception as e:
+        status_text.error(f"❌ Error: {e}")
+        log(f"ERROR: {e}")
+        log(traceback.format_exc())
+    finally:
+        st.session_state["bench_running"] = False
+
+
+def run_indexing_with_ui(config: Dict[str, Any]):
+    st.divider()
+
+    progress_bar = st.progress(0.0)
+    status_text = st.empty()
+    output_area = st.empty()
+
+    status_text.info("🚀 Starting indexing...")
+    output_lines = []
+
+    def log(msg):
+        output_lines.append(msg)
+        output_area.code("\n".join(output_lines[-50:]), language="text")
+
+    try:
+        url, api_key = get_qdrant_credentials()
+        if not url:
+            st.error("QDRANT_URL not configured")
+            return
+
+        datasets = config.get("datasets", [])
+        collection = config["collection"]
+        model = config.get("model", "vidore/colpali-v1.3")
+        recreate = config.get("recreate", False)
+        torch_dtype = config.get("torch_dtype", "float16")
+        qdrant_vector_dtype = config.get("qdrant_vector_dtype", "float16")
+        prefer_grpc = config.get("prefer_grpc", True)
+        batch_size = config.get("batch_size", 4)
+
+        log(f"[Index] Model: {model}")
+        log(f"[Index] Collection: {collection}")
+        log(f"[Index] Datasets: {datasets}")
+        status_text.info("📦 Loading embedder...")
+
+        embedder = VisualEmbedder(model_name=model, torch_dtype=torch_dtype)
+        log("[Index] Embedder loaded")
+
+        status_text.info("🔌 Connecting to Qdrant...")
+        indexer = QdrantIndexer(
+            url=url,
+            api_key=api_key,
+            collection_name=collection,
+            prefer_grpc=prefer_grpc,
+            vector_datatype=qdrant_vector_dtype,
+        )
+        log("[Index] Connected to Qdrant")
+
+        status_text.info("📦 Creating collection...")
+        indexer.create_collection(force_recreate=recreate)
+        indexer.create_payload_indexes(
+            fields=[
+                {"field": "dataset", "type": "keyword"},
+                {"field": "doc_id", "type": "keyword"},
+            ]
+        )
+        log(f"[Index] Collection '{collection}' ready")
+
+        total_uploaded = 0
+
+        for ds_name in datasets:
+            status_text.info(f"📚 Loading dataset: {ds_name}")
+            corpus, queries, qrels = load_vidore_beir_dataset(ds_name)
+            log(f"[Index] Loaded {ds_name}: {len(corpus)} documents")
+
+            for i in range(0, len(corpus), batch_size):
+                batch = corpus[i : i + batch_size]
+                images = [doc.image for doc in batch if hasattr(doc, "image") and doc.image]
+
+                if not images:
+                    continue
+
+                status_text.info(f"🎨 Embedding batch {i//batch_size + 1}...")
+                embeddings = embedder.embed_images(images)
+
+                points = []
+                for j, (doc, emb) in enumerate(zip(batch, embeddings)):
+                    doc_id = doc.doc_id if hasattr(doc, "doc_id") else str(i + j)
+                    emb_np = emb.cpu().numpy() if hasattr(emb, "cpu") else np.array(emb)
+                    tile_pooled = emb_np.reshape(-1, 4, emb_np.shape[-1]).mean(axis=1)
+                    global_pooled = emb_np.mean(axis=0)
+
+                    points.append(
+                        {
+                            "id": f"{ds_name}_{doc_id}".replace("/", "_"),
+                            "visual_embedding": emb_np,
+                            "tile_pooled_embedding": tile_pooled,
+                            "experimental_pooled_embedding": tile_pooled,
+                            "global_pooled_embedding": global_pooled,
+                            "metadata": {"dataset": ds_name, "doc_id": doc_id},
+                        }
+                    )
+
+                indexer.upload_batch(points)
+                total_uploaded += len(points)
+
+                progress = (i + len(batch)) / len(corpus)
+                progress_bar.progress(progress)
+                log(f"[Index] Uploaded {total_uploaded} points")
+
+        progress_bar.progress(1.0)
+        status_text.success(f"✅ Indexing complete! {total_uploaded} documents indexed.")
+
+    except Exception as e:
+        status_text.error(f"❌ Error: {e}")
+        log(f"ERROR: {e}")
+        log(traceback.format_exc())
+
+
+def render_benchmark_results():
+    st.markdown("##### Load Results")
+
+    available = get_available_results()
+
+    if not available:
+        st.info("No results found")
+        return
+
+    default_select = []
+    if st.session_state.get("auto_select_result"):
+        auto = st.session_state.pop("auto_select_result")
+        if auto in [str(p) for p in available]:
+            default_select = [auto]
+
+    selected = st.multiselect(
+        "Result files",
+        options=[str(p) for p in available],
+        format_func=lambda x: Path(x).name[:60],
+        default=default_select,
+        key="br_files",
+    )
+
+    for path in selected:
+        data = load_results_file(Path(path))
+        if data:
+            render_result_card(data, Path(path).name)
+
+
+def render_result_card(data: Dict[str, Any], filename: str):
+    with st.expander(f"📊 {filename[:50]}", expanded=True):
+        c1, c2, c3, c4 = st.columns(4)
+        c1.metric("Model", (data.get("model") or "?").split("/")[-1])
+        c2.metric("Mode", data.get("mode", "?"))
+        c3.metric("Top K", data.get("top_k", "?"))
+        c4.metric("Time", f"{data.get('eval_wall_time_s', 0):.0f}s")
+
+        metrics = data.get("metrics_by_dataset", {})
+        if not metrics:
+            st.warning("No metrics data")
+            return
+
+        rows = []
+        for ds, m in metrics.items():
+            rows.append(
+                {
+                    "Dataset": ds.split("/")[-1].replace("_v2", ""),
+                    "NDCG@5": m.get("ndcg@5", 0),
+                    "NDCG@10": m.get("ndcg@10", 0),
+                    "Recall@5": m.get("recall@5", 0),
+                    "Recall@10": m.get("recall@10", 0),
+                    "MRR@10": m.get("mrr@10", 0),
+                    "Latency": m.get("avg_latency_ms", 0),
+                    "QPS": m.get("qps", 0),
+                }
+            )
+
+        df = pd.DataFrame(rows)
+
+        st.dataframe(
+            df.style.format(
+                {
+                    "NDCG@5": "{:.4f}",
+                    "NDCG@10": "{:.4f}",
+                    "Recall@5": "{:.4f}",
+                    "Recall@10": "{:.4f}",
+                    "MRR@10": "{:.4f}",
+                    "Latency": "{:.1f}",
+                    "QPS": "{:.2f}",
+                }
+            ),
+            hide_index=True,
+            use_container_width=True,
+        )
+
+        chart_data = []
+        for ds, m in metrics.items():
+            ds_short = ds.split("/")[-1].replace("_v2", "").replace("_", " ").title()
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "NDCG@10", "Value": m.get("ndcg@10", 0)}
+            )
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "Recall@10", "Value": m.get("recall@10", 0)}
+            )
+            chart_data.append(
+                {"Dataset": ds_short, "Metric": "MRR@10", "Value": m.get("mrr@10", 0)}
+            )
+
+        chart_df = pd.DataFrame(chart_data)
+
+        chart = (
+            alt.Chart(chart_df)
+            .mark_bar()
+            .encode(
+                x=alt.X("Dataset:N", title=None),
+                y=alt.Y("Value:Q", scale=alt.Scale(domain=[0, 1]), title="Score"),
+                color=alt.Color("Metric:N", scale=alt.Scale(scheme="tableau10")),
+                xOffset="Metric:N",
+                tooltip=["Dataset", "Metric", alt.Tooltip("Value:Q", format=".4f")],
+            )
+            .properties(height=300, title="Metrics by Dataset")
+        )
+
+        st.altair_chart(chart, use_container_width=True)
+
+        latency_data = [
+            {
+                "Dataset": ds.split("/")[-1].replace("_v2", ""),
+                "Latency (ms)": m.get("avg_latency_ms", 0),
+                "QPS": m.get("qps", 0),
+            }
+            for ds, m in metrics.items()
+        ]
+        latency_df = pd.DataFrame(latency_data)
+
+        c1, c2 = st.columns(2)
+        with c1:
+            lat_chart = (
+                alt.Chart(latency_df)
+                .mark_bar(color="#ff6b6b")
+                .encode(
+                    x=alt.X("Dataset:N"),
+                    y=alt.Y("Latency (ms):Q"),
+                    tooltip=["Dataset", alt.Tooltip("Latency (ms):Q", format=".1f")],
+                )
+                .properties(height=200, title="Avg Latency")
+            )
+            st.altair_chart(lat_chart, use_container_width=True)
+
+        with c2:
+            qps_chart = (
+                alt.Chart(latency_df)
+                .mark_bar(color="#4ecdc4")
+                .encode(
+                    x=alt.X("Dataset:N"),
+                    y=alt.Y("QPS:Q"),
+                    tooltip=["Dataset", alt.Tooltip("QPS:Q", format=".2f")],
+                )
+                .properties(height=200, title="QPS (Queries/sec)")
+            )
+            st.altair_chart(qps_chart, use_container_width=True)
+
+
+def main():
+    render_header()
+    render_sidebar()
+
+    tab_upload, tab_playground, tab_benchmark = st.tabs(
+        ["📤 Upload", "🎮 Playground", "📊 Benchmarking"]
+    )
+
+    with tab_upload:
+        render_upload_tab()
+
+    with tab_playground:
+        render_playground_tab()
+
+    with tab_benchmark:
+        render_benchmark_tab()
+
+
+if __name__ == "__main__":
+    main()

examples/COMMANDS.md

@@ -57,7 +57,7 @@ python -m benchmarks.vidore_tatdqa_test.sweep_eval \
   --collection vidore_tatdqa_test \
   --prefer-grpc \
   --mode two_stage \
-  --stage1-mode tokens_vs_tiles \
+  --stage1-mode tokens_vs_standard_pooling \
   --prefetch-ks 20,50,100,200,400 \
   --torch-dtype auto \
   --query-batch-size 32 \
@@ -80,4 +80,98 @@ python -m benchmarks.vidore_tatdqa_test.sweep_eval \
   --out-dir results/sweeps
 ```
 
+---
+
+# ViDoRe v2 BEIR datasets (Qdrant) — commands
+
+This section indexes the **3 ViDoRe v2** datasets used in the demo UI:
+
+- `vidore/esg_reports_v2`
+- `vidore/biomedical_lectures_v2`
+- `vidore/economics_reports_v2`
+
+We use **`vidore/colqwen2.5-v0.2`**, **no cropping**, **no Cloudinary**, **gRPC**, and **float32** for both compute and stored vectors.
+
+## Environment
+
+```bash
+export QDRANT_URL="https://YOUR_QDRANT_HOST:6333"
+export QDRANT_API_KEY="YOUR_KEY"  # optional for local Qdrant
+```
+
+Optional (recommended on machines with small disks):
+
+```bash
+export HF_HOME="$PWD/.cache/huggingface"
+export TRANSFORMERS_CACHE="$PWD/.cache/huggingface"
+```
+
+## Index only (no evaluation)
+
+```bash
+python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir \
+  --datasets \
+    vidore/esg_reports_v2 \
+    vidore/biomedical_lectures_v2 \
+    vidore/economics_reports_v2 \
+  --collection vidore_v2__colqwen25_fp32 \
+  --model vidore/colqwen2.5-v0.2 \
+  --index \
+  --recreate \
+  --indexing-threshold 0 \
+  --full-scan-threshold 0 \
+  --prefer-grpc \
+  --torch-dtype float32 \
+  --qdrant-vector-dtype float32 \
+  --batch-size 1 \
+  --upload-batch-size 4 \
+  --upload-workers 0 \
+  --no-cloudinary \
+  --no-eval
+```
+
+Notes:
+- **`--batch-size 1`** is the safest starting point on Apple Silicon (MPS). Increase cautiously if stable.
+- This does **not** enable cropping (we do **not** pass `--crop-empty`).
+
+## Evaluate later (optional)
+
+Single-stage full MaxSim:
+
+```bash
+python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir \
+  --datasets \
+    vidore/esg_reports_v2 \
+    vidore/biomedical_lectures_v2 \
+    vidore/economics_reports_v2 \
+  --collection vidore_v2__colqwen25_fp32 \
+  --model vidore/colqwen2.5-v0.2 \
+  --prefer-grpc \
+  --torch-dtype float32 \
+  --qdrant-vector-dtype float32 \
+  --mode single_full \
+  --top-k 100 \
+  --evaluation-scope per_dataset
+```
+
+Two-stage (prefetch + rerank):
+
+```bash
+python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir \
+  --datasets \
+    vidore/esg_reports_v2 \
+    vidore/biomedical_lectures_v2 \
+    vidore/economics_reports_v2 \
+  --collection vidore_v2__colqwen25_fp32 \
+  --model vidore/colqwen2.5-v0.2 \
+  --prefer-grpc \
+  --torch-dtype float32 \
+  --qdrant-vector-dtype float32 \
+  --mode two_stage \
+  --stage1-mode tokens_vs_experimental_pooling \
+  --prefetch-k 200 \
+  --top-k 100 \
+  --evaluation-scope per_dataset
+```
+
 

examples/process_pdfs.py

@@ -10,18 +10,18 @@ This example demonstrates the full pipeline:
 
 Usage:
     python examples/process_pdfs.py --reports-dir /path/to/pdfs
-    
+
     # With metadata mapping
     python examples/process_pdfs.py --reports-dir /path/to/pdfs --metadata-file metadata.json
-    
+
     # Without Cloudinary (local embeddings only)
     python examples/process_pdfs.py --reports-dir /path/to/pdfs --no-cloudinary
 """
 
-import os
-import sys
 import argparse
 import logging
+import os
+import sys
 from pathlib import Path
 
 from dotenv import load_dotenv
@@ -29,12 +29,11 @@ from dotenv import load_dotenv
 # Add parent to path for development
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from visual_rag import VisualEmbedder, QdrantIndexer, CloudinaryUploader, load_config
-from visual_rag.indexing.pipeline import ProcessingPipeline
+from visual_rag import CloudinaryUploader, QdrantIndexer, VisualEmbedder, load_config  # noqa: E402
+from visual_rag.indexing.pipeline import ProcessingPipeline  # noqa: E402
 
 logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 )
 logger = logging.getLogger(__name__)
 
@@ -42,97 +41,83 @@ logger = logging.getLogger(__name__)
 def main():
     parser = argparse.ArgumentParser(description="Process PDFs with Visual RAG Toolkit")
     parser.add_argument(
-        "--reports-dir", type=str, required=True,
-        help="Directory containing PDF files"
-    )
-    parser.add_argument(
-        "--metadata-file", type=str,
-        help="JSON file with filename → metadata mapping (optional)"
-    )
-    parser.add_argument(
-        "--config", type=str, default="config.yaml",
-        help="Configuration file path"
-    )
-    parser.add_argument(
-        "--collection", type=str,
-        help="Qdrant collection name (overrides config)"
+        "--reports-dir", type=str, required=True, help="Directory containing PDF files"
     )
     parser.add_argument(
-        "--model", type=str,
-        help="Model name (overrides config)"
+        "--metadata-file", type=str, help="JSON file with filename → metadata mapping (optional)"
     )
+    parser.add_argument("--config", type=str, default="config.yaml", help="Configuration file path")
+    parser.add_argument("--collection", type=str, help="Qdrant collection name (overrides config)")
+    parser.add_argument("--model", type=str, help="Model name (overrides config)")
+    parser.add_argument("--no-cloudinary", action="store_true", help="Skip Cloudinary uploads")
     parser.add_argument(
-        "--no-cloudinary", action="store_true",
-        help="Skip Cloudinary uploads"
+        "--no-qdrant", action="store_true", help="Skip Qdrant uploads (just generate embeddings)"
     )
     parser.add_argument(
-        "--no-qdrant", action="store_true",
-        help="Skip Qdrant uploads (just generate embeddings)"
+        "--skip-existing",
+        action="store_true",
+        default=True,
+        help="Skip pages that already exist in Qdrant (default: True)",
     )
-    parser.add_argument(
-        "--skip-existing", action="store_true", default=True,
-        help="Skip pages that already exist in Qdrant (default: True)"
-    )
-    parser.add_argument(
-        "--force", action="store_true",
-        help="Process all pages even if they exist"
-    )
-    
+    parser.add_argument("--force", action="store_true", help="Process all pages even if they exist")
+
     args = parser.parse_args()
-    
+
     # Load environment variables
     load_dotenv()
-    
+
     # Load configuration
     config = load_config(args.config)
-    
+
     # Get PDFs
     reports_dir = Path(args.reports_dir)
     if not reports_dir.exists():
         logger.error(f"Reports directory not found: {reports_dir}")
         sys.exit(1)
-    
+
     pdf_paths = sorted(reports_dir.glob("*.pdf")) + sorted(reports_dir.glob("*.PDF"))
     if not pdf_paths:
         logger.error(f"No PDF files found in: {reports_dir}")
         sys.exit(1)
-    
+
     logger.info(f"📁 Found {len(pdf_paths)} PDF files in {reports_dir}")
-    
+
     # Load metadata mapping if provided
     metadata_mapping = {}
     if args.metadata_file:
         metadata_mapping = ProcessingPipeline.load_metadata_mapping(Path(args.metadata_file))
-    
+
     # Get settings
     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")
-    
+    collection_name = args.collection or config.get("qdrant", {}).get(
+        "collection_name", "visual_documents"
+    )
+
     # Initialize embedder
     logger.info(f"🤖 Initializing embedder: {model_name}")
     embedder = VisualEmbedder(model_name=model_name)
-    
+
     # Initialize Qdrant indexer (if not skipped)
     indexer = None
     if not args.no_qdrant:
         qdrant_url = os.getenv("QDRANT_URL")
         qdrant_api_key = os.getenv("QDRANT_API_KEY")
-        
+
         if not qdrant_url:
             logger.error("QDRANT_URL environment variable not set")
             sys.exit(1)
-        
+
         logger.info(f"🔌 Connecting to Qdrant: {qdrant_url}")
         indexer = QdrantIndexer(
             url=qdrant_url,
             api_key=qdrant_api_key,
             collection_name=collection_name,
         )
-        
+
         # Create collection if needed
         indexer.create_collection()
         indexer.create_payload_indexes()
-    
+
     # Initialize Cloudinary uploader (if not skipped)
     cloudinary_uploader = None
     if not args.no_cloudinary:
@@ -143,7 +128,7 @@ def main():
         except ValueError as e:
             logger.warning(f"Cloudinary not configured: {e}")
             logger.warning("Continuing without Cloudinary uploads")
-    
+
     # Create pipeline
     pipeline = ProcessingPipeline(
         embedder=embedder,
@@ -152,39 +137,39 @@ def main():
         metadata_mapping=metadata_mapping,
         config=config,
     )
-    
+
     # Process PDFs
     total_uploaded = 0
     total_skipped = 0
     total_failed = 0
-    
+
     skip_existing = args.skip_existing and not args.force
-    
+
     for pdf_idx, pdf_path in enumerate(pdf_paths, 1):
         logger.info(f"\n{'='*60}")
         logger.info(f"📄 [{pdf_idx}/{len(pdf_paths)}] {pdf_path.name}")
         logger.info(f"{'='*60}")
-        
+
         result = pipeline.process_pdf(
             pdf_path,
             skip_existing=skip_existing,
             upload_to_cloudinary=(not args.no_cloudinary),
             upload_to_qdrant=(not args.no_qdrant),
         )
-        
+
         total_uploaded += result["uploaded"]
         total_skipped += result["skipped"]
         total_failed += result["failed"]
-    
+
     # Summary
     logger.info(f"\n{'='*60}")
-    logger.info(f"📊 SUMMARY")
+    logger.info("📊 SUMMARY")
     logger.info(f"{'='*60}")
     logger.info(f"   Total PDFs: {len(pdf_paths)}")
     logger.info(f"   Uploaded: {total_uploaded}")
     logger.info(f"   Skipped: {total_skipped}")
     logger.info(f"   Failed: {total_failed}")
-    
+
     if indexer:
         info = indexer.get_collection_info()
         if info:
@@ -193,10 +178,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-
-
-
-
-

examples/search_demo.py

@@ -9,18 +9,18 @@ This example demonstrates:
 
 Usage:
     python examples/search_demo.py --query "What is the budget allocation?"
-    
+
     # With filters
     python examples/search_demo.py --query "budget" --year 2023 --source "Local Government"
-    
+
     # With saliency maps
     python examples/search_demo.py --query "budget" --saliency
 """
 
-import os
-import sys
 import argparse
 import logging
+import os
+import sys
 from pathlib import Path
 
 from dotenv import load_dotenv
@@ -29,9 +29,9 @@ from qdrant_client import QdrantClient
 # Add parent to path for development
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from visual_rag import VisualEmbedder
-from visual_rag.retrieval.two_stage import TwoStageRetriever
-from visual_rag.visualization import visualize_search_results, generate_saliency_map
+from visual_rag import VisualEmbedder  # noqa: E402
+from visual_rag.retrieval.two_stage import TwoStageRetriever  # noqa: E402
+from visual_rag.visualization import visualize_search_results  # noqa: E402
 
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -39,82 +39,58 @@ logger = logging.getLogger(__name__)
 
 def main():
     parser = argparse.ArgumentParser(description="Search with Visual RAG Toolkit")
+    parser.add_argument("--query", type=str, required=True, help="Search query")
     parser.add_argument(
-        "--query", type=str, required=True,
-        help="Search query"
-    )
-    parser.add_argument(
-        "--collection", type=str, default="visual_documents",
-        help="Qdrant collection name"
-    )
-    parser.add_argument(
-        "--model", type=str, default="vidore/colSmol-500M",
-        help="Model name"
-    )
-    parser.add_argument(
-        "--top-k", type=int, default=10,
-        help="Number of results"
-    )
-    parser.add_argument(
-        "--prefetch-k", type=int, default=200,
-        help="Candidates for two-stage retrieval"
+        "--collection", type=str, default="visual_documents", help="Qdrant collection name"
     )
+    parser.add_argument("--model", type=str, default="vidore/colSmol-500M", help="Model name")
+    parser.add_argument("--top-k", type=int, default=10, help="Number of results")
     parser.add_argument(
-        "--year", type=int,
-        help="Filter by year"
+        "--prefetch-k", type=int, default=200, help="Candidates for two-stage retrieval"
     )
+    parser.add_argument("--year", type=int, help="Filter by year")
+    parser.add_argument("--source", type=str, help="Filter by source")
+    parser.add_argument("--district", type=str, help="Filter by district")
     parser.add_argument(
-        "--source", type=str,
-        help="Filter by source"
+        "--saliency", action="store_true", help="Generate saliency maps for results"
     )
-    parser.add_argument(
-        "--district", type=str,
-        help="Filter by district"
-    )
-    parser.add_argument(
-        "--saliency", action="store_true",
-        help="Generate saliency maps for results"
-    )
-    parser.add_argument(
-        "--output", type=str,
-        help="Output path for visualization"
-    )
-    
+    parser.add_argument("--output", type=str, help="Output path for visualization")
+
     args = parser.parse_args()
-    
+
     # Load environment
     load_dotenv()
-    
+
     qdrant_url = os.getenv("QDRANT_URL")
     qdrant_api_key = os.getenv("QDRANT_API_KEY")
-    
+
     if not qdrant_url:
         logger.error("QDRANT_URL not set")
         sys.exit(1)
-    
+
     # Initialize components
     logger.info(f"🤖 Loading model: {args.model}")
     embedder = VisualEmbedder(model_name=args.model)
-    
+
     logger.info(f"🔌 Connecting to Qdrant: {qdrant_url}")
     client = QdrantClient(url=qdrant_url, api_key=qdrant_api_key)
-    
+
     retriever = TwoStageRetriever(
         qdrant_client=client,
         collection_name=args.collection,
     )
-    
+
     # Embed query
     logger.info(f"🔍 Query: {args.query}")
     query_embedding = embedder.embed_query(args.query)
-    
+
     # Build filter
     filter_obj = retriever.build_filter(
         year=args.year,
         source=args.source,
         district=args.district,
     )
-    
+
     # Search
     results = retriever.search(
         query_embedding=query_embedding.numpy(),
@@ -122,31 +98,31 @@ def main():
         prefetch_k=args.prefetch_k,
         filter_obj=filter_obj,
     )
-    
+
     # Display results
     logger.info(f"\n📊 Results ({len(results)}):")
     for i, result in enumerate(results, 1):
         payload = result.get("payload", {})
         score = result.get("score_final", result.get("score_stage1", 0))
-        
+
         filename = payload.get("filename", "N/A")
         page_num = payload.get("page_number", "N/A")
         year = payload.get("year", "N/A")
         source = payload.get("source", "N/A")
-        
+
         logger.info(f"  {i}. {filename} p.{page_num}")
         logger.info(f"     Score: {score:.4f} | Year: {year} | Source: {source}")
-        
+
         # Show text snippet
         text = payload.get("text", "")
         if text:
             snippet = text[:200].replace("\n", " ")
             logger.info(f"     Text: {snippet}...")
-    
+
     # Visualize results
     if args.output or args.saliency:
         output_path = args.output or "search_results.png"
-        
+
         logger.info(f"\n🎨 Generating visualization: {output_path}")
         visualize_search_results(
             query=args.query,
@@ -158,10 +134,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-
-
-
-
-
-

scripts/colqwen25_probe.py

@@ -0,0 +1,70 @@
+"""
+Probe script for vidore/colqwen2.5-v0.2 embedding layout.
+
+Usage:
+  python scripts/colqwen25_probe.py --model vidore/colqwen2.5-v0.2 --device cuda:0
+
+Notes:
+- ColQwen2.5 requires colpali-engine>=0.3.7 and transformers>=4.45.0
+  (the model card recommends installing from source).
+- This script prints embedding shapes + token_info (grid_h/grid_w when available),
+  and runs mean/experimental pooling to validate compatibility with the pipeline.
+"""
+
+from __future__ import annotations
+
+import argparse
+
+from PIL import Image
+
+from visual_rag import VisualEmbedder
+
+
+def main() -> None:
+    p = argparse.ArgumentParser()
+    p.add_argument("--model", default="vidore/colqwen2.5-v0.2")
+    p.add_argument("--device", default="cpu")
+    p.add_argument("--dtype", default="bfloat16", choices=["bfloat16", "float16", "float32"])
+    args = p.parse_args()
+
+    img = Image.new("RGB", (1024, 768), color="white")
+
+    embedder = VisualEmbedder(model_name=args.model, device=args.device, torch_dtype=args.dtype)
+    embs, infos = embedder.embed_images([img], return_token_info=True, show_progress=False)
+
+    emb = embs[0]
+    info = infos[0]
+    print("Model:", args.model)
+    print("Full embedding:", tuple(emb.shape), "dtype:", emb.dtype)
+    print(
+        "token_info:",
+        {
+            k: info.get(k)
+            for k in [
+                "num_visual_tokens",
+                "grid_t",
+                "grid_h",
+                "grid_w",
+                "n_rows",
+                "n_cols",
+                "num_tiles",
+            ]
+        },
+    )
+
+    visual = embedder.extract_visual_embedding(emb, info)
+    print("Visual embedding:", tuple(visual.shape), "dtype:", visual.dtype)
+
+    mean_pool = embedder.mean_pool_visual_embedding(visual, info, target_vectors=32)
+    exp_pool = embedder.experimental_pool_visual_embedding(
+        visual, info, target_vectors=32, mean_pool=mean_pool
+    )
+    global_pool = embedder.global_pool_from_mean_pool(mean_pool)
+
+    print("mean_pool:", tuple(mean_pool.shape))
+    print("exp_pool:", tuple(exp_pool.shape))
+    print("global_pool:", tuple(global_pool.shape))
+
+
+if __name__ == "__main__":
+    main()

scripts/compare_eval_scopes.py

@@ -0,0 +1,309 @@
+"""
+Run the same benchmark twice (union vs per_dataset) and print full reports + deltas.
+
+This is designed to answer: "How much do distractors (union scope) hurt vs per-dataset filtering?"
+
+It runs:
+  python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir ... --evaluation-scope union
+  python -m benchmarks.vidore_beir_qdrant.run_qdrant_beir ... --evaluation-scope per_dataset
+
+Then prints, per dataset:
+  - full metrics dict for union
+  - full metrics dict for per_dataset
+  - delta = per_dataset - union (for numeric metrics)
+
+Usage example:
+  python scripts/compare_eval_scopes.py \\
+    --datasets vidore/esg_reports_v2 vidore/biomedical_lectures_v2 vidore/economics_reports_v2 \\
+    --collection vidore_beir_v2_3ds__colpali_v1_3__nocrop__union \\
+    --model vidore/colpali-v1.3 \\
+    --mode single_full \\
+    --top-k 100
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import subprocess
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+
+def _now_tag() -> str:
+    return datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+
+
+def _as_number(x: Any) -> Optional[float]:
+    if x is None:
+        return None
+    if isinstance(x, bool):
+        return float(int(x))
+    if isinstance(x, (int, float)):
+        return float(x)
+    return None
+
+
+def _delta_metrics(per_ds: Dict[str, Any], union: Dict[str, Any]) -> Dict[str, Any]:
+    out: Dict[str, Any] = {}
+    keys = set(per_ds.keys()) | set(union.keys())
+    for k in sorted(keys):
+        a = _as_number(per_ds.get(k))
+        b = _as_number(union.get(k))
+        if a is not None and b is not None:
+            out[k] = a - b
+        else:
+            # keep non-numerics as a tuple when present in either
+            if k in per_ds or k in union:
+                out[k] = {"per_dataset": per_ds.get(k), "union": union.get(k)}
+    return out
+
+
+def _load_metrics_by_dataset(path: Path) -> Dict[str, Dict[str, Any]]:
+    obj = json.loads(path.read_text())
+    mbd = obj.get("metrics_by_dataset") or {}
+    if not isinstance(mbd, dict):
+        return {}
+    # ensure nested dicts
+    out: Dict[str, Dict[str, Any]] = {}
+    for k, v in mbd.items():
+        if isinstance(v, dict):
+            out[str(k)] = v
+    return out
+
+
+def _run_once(
+    *,
+    datasets: List[str],
+    collection: str,
+    model: str,
+    mode: str,
+    top_k: int,
+    stage1_mode: Optional[str],
+    prefetch_k: Optional[int],
+    stage1_k: Optional[int],
+    stage2_k: Optional[int],
+    torch_dtype: str,
+    qdrant_vector_dtype: str,
+    prefer_grpc: bool,
+    max_queries: int,
+    evaluation_scope: str,
+    qdrant_timeout: int,
+    qdrant_retries: int,
+    qdrant_retry_sleep: float,
+    extra_args: List[str],
+    out_path: Path,
+) -> None:
+    cmd: List[str] = [
+        sys.executable,
+        "-m",
+        "benchmarks.vidore_beir_qdrant.run_qdrant_beir",
+        "--datasets",
+        *datasets,
+        "--collection",
+        collection,
+        "--model",
+        model,
+        "--mode",
+        mode,
+        "--top-k",
+        str(int(top_k)),
+        "--evaluation-scope",
+        str(evaluation_scope),
+        "--torch-dtype",
+        torch_dtype,
+        "--qdrant-vector-dtype",
+        qdrant_vector_dtype,
+        "--qdrant-timeout",
+        str(int(qdrant_timeout)),
+        "--qdrant-retries",
+        str(int(qdrant_retries)),
+        "--qdrant-retry-sleep",
+        str(float(qdrant_retry_sleep)),
+        "--max-queries",
+        str(int(max_queries)),
+        "--output",
+        str(out_path),
+    ]
+
+    if not prefer_grpc:
+        cmd.append("--no-prefer-grpc")
+    else:
+        cmd.append("--prefer-grpc")
+
+    if str(mode) == "two_stage":
+        if stage1_mode:
+            cmd += ["--stage1-mode", str(stage1_mode)]
+        if prefetch_k is not None:
+            cmd += ["--prefetch-k", str(int(prefetch_k))]
+    if str(mode) == "three_stage":
+        if stage1_k is not None:
+            cmd += ["--stage1-k", str(int(stage1_k))]
+        if stage2_k is not None:
+            cmd += ["--stage2-k", str(int(stage2_k))]
+
+    cmd += list(extra_args or [])
+
+    env = os.environ.copy()
+    env.setdefault("HF_HUB_DISABLE_XET", "1")  # avoid xet crashes in some environments
+
+    print("\n" + "=" * 90)
+    print(f"RUN scope={evaluation_scope}")
+    print(" ".join(cmd))
+    print("=" * 90)
+    sys.stdout.flush()
+
+    subprocess.run(cmd, check=True, env=env)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--datasets", nargs="+", required=True)
+    ap.add_argument("--collection", required=True)
+    ap.add_argument("--model", required=True)
+    ap.add_argument(
+        "--mode", default="single_full", choices=["single_full", "two_stage", "three_stage"]
+    )
+    ap.add_argument("--top-k", type=int, default=100)
+    ap.add_argument(
+        "--stage1-mode",
+        default="",
+        help="two_stage stage1 mode (e.g. tokens_vs_experimental_pooling or tokens_vs_standard_pooling)",
+    )
+    ap.add_argument("--prefetch-k", type=int, default=256)
+    ap.add_argument("--stage1-k", type=int, default=1000)
+    ap.add_argument("--stage2-k", type=int, default=300)
+    ap.add_argument(
+        "--torch-dtype", default="auto", choices=["auto", "float32", "float16", "bfloat16"]
+    )
+    ap.add_argument("--qdrant-vector-dtype", default="float16", choices=["float16", "float32"])
+    ap.add_argument("--prefer-grpc", action="store_true", default=False)
+    ap.add_argument("--max-queries", type=int, default=0)
+    ap.add_argument("--qdrant-timeout", type=int, default=120)
+    ap.add_argument("--qdrant-retries", type=int, default=3)
+    ap.add_argument("--qdrant-retry-sleep", type=float, default=0.5)
+    ap.add_argument(
+        "--out-dir",
+        default="results/scope_comparisons",
+        help="Directory to write the two raw JSON reports + the merged report",
+    )
+    ap.add_argument(
+        "--extra-arg",
+        action="append",
+        default=[],
+        help="Pass-through extra args to run_qdrant_beir (repeatable), e.g. --extra-arg --crop-empty",
+    )
+    args = ap.parse_args()
+
+    datasets = [str(x) for x in args.datasets]
+    stage1_mode = str(args.stage1_mode).strip() or None
+
+    out_dir = Path(str(args.out_dir))
+    out_dir.mkdir(parents=True, exist_ok=True)
+    tag = _now_tag()
+
+    base = f"{tag}__{Path(args.collection).name}__{Path(args.model).name}__{args.mode}"
+    union_path = out_dir / f"{base}__scope_union.json"
+    per_path = out_dir / f"{base}__scope_per_dataset.json"
+    merged_path = out_dir / f"{base}__scope_compare.json"
+
+    _run_once(
+        datasets=datasets,
+        collection=str(args.collection),
+        model=str(args.model),
+        mode=str(args.mode),
+        top_k=int(args.top_k),
+        stage1_mode=stage1_mode,
+        prefetch_k=int(args.prefetch_k),
+        stage1_k=int(args.stage1_k),
+        stage2_k=int(args.stage2_k),
+        torch_dtype=str(args.torch_dtype),
+        qdrant_vector_dtype=str(args.qdrant_vector_dtype),
+        prefer_grpc=bool(args.prefer_grpc),
+        max_queries=int(args.max_queries),
+        evaluation_scope="union",
+        qdrant_timeout=int(args.qdrant_timeout),
+        qdrant_retries=int(args.qdrant_retries),
+        qdrant_retry_sleep=float(args.qdrant_retry_sleep),
+        extra_args=list(args.extra_arg or []),
+        out_path=union_path,
+    )
+    _run_once(
+        datasets=datasets,
+        collection=str(args.collection),
+        model=str(args.model),
+        mode=str(args.mode),
+        top_k=int(args.top_k),
+        stage1_mode=stage1_mode,
+        prefetch_k=int(args.prefetch_k),
+        stage1_k=int(args.stage1_k),
+        stage2_k=int(args.stage2_k),
+        torch_dtype=str(args.torch_dtype),
+        qdrant_vector_dtype=str(args.qdrant_vector_dtype),
+        prefer_grpc=bool(args.prefer_grpc),
+        max_queries=int(args.max_queries),
+        evaluation_scope="per_dataset",
+        qdrant_timeout=int(args.qdrant_timeout),
+        qdrant_retries=int(args.qdrant_retries),
+        qdrant_retry_sleep=float(args.qdrant_retry_sleep),
+        extra_args=list(args.extra_arg or []),
+        out_path=per_path,
+    )
+
+    union_mbd = _load_metrics_by_dataset(union_path)
+    per_mbd = _load_metrics_by_dataset(per_path)
+
+    all_ds = sorted(set(union_mbd.keys()) | set(per_mbd.keys()))
+    comparison: Dict[str, Any] = {
+        "meta": {
+            "datasets": datasets,
+            "collection": str(args.collection),
+            "model": str(args.model),
+            "mode": str(args.mode),
+            "top_k": int(args.top_k),
+            "stage1_mode": stage1_mode,
+            "prefetch_k": int(args.prefetch_k) if str(args.mode) == "two_stage" else None,
+            "stage1_k": int(args.stage1_k) if str(args.mode) == "three_stage" else None,
+            "stage2_k": int(args.stage2_k) if str(args.mode) == "three_stage" else None,
+            "torch_dtype": str(args.torch_dtype),
+            "qdrant_vector_dtype": str(args.qdrant_vector_dtype),
+            "prefer_grpc": bool(args.prefer_grpc),
+            "max_queries": int(args.max_queries),
+            "union_report": str(union_path),
+            "per_dataset_report": str(per_path),
+        },
+        "by_dataset": {},
+    }
+
+    print("\n" + "#" * 90)
+    print("SCOPE COMPARISON (per_dataset − union)")
+    print("#" * 90)
+    for ds in all_ds:
+        u = union_mbd.get(ds, {})
+        p = per_mbd.get(ds, {})
+        d = _delta_metrics(p, u)
+        comparison["by_dataset"][ds] = {
+            "union": u,
+            "per_dataset": p,
+            "delta": d,
+        }
+        print("\n" + "-" * 90)
+        print(ds)
+        print("-" * 90)
+        print("UNION:")
+        print(json.dumps(u, indent=2, sort_keys=True))
+        print("\nPER_DATASET:")
+        print(json.dumps(p, indent=2, sort_keys=True))
+        print("\nDELTA (per_dataset - union):")
+        print(json.dumps(d, indent=2, sort_keys=True))
+        sys.stdout.flush()
+
+    merged_path.write_text(json.dumps(comparison, indent=2, sort_keys=True))
+    print("\nWrote merged comparison:", merged_path)
+
+
+if __name__ == "__main__":
+    main()

scripts/compare_models_sample_queries.py

@@ -0,0 +1,290 @@
+"""
+Compare retrieval quality across two model+collection pairs on the same dataset queries.
+
+This is a read-only diagnostic:
+- Loads BEIR dataset (queries + qrels)
+- Remaps qrels doc_ids -> Qdrant point IDs for each collection
+- Runs retrieval for a sample of queries
+- Computes simple hit-rate statistics + per-query best-rank
+- Writes a JSON report under results/model_compare/
+
+Example:
+  python scripts/compare_models_sample_queries.py \\
+    --dataset vidore/esg_reports_v2 \\
+    --top-k 100 \\
+    --max-queries 50
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import hashlib
+import json
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+from qdrant_client.http import models as qm
+
+from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+from visual_rag import VisualEmbedder
+from visual_rag.retrieval import MultiVectorRetriever
+
+
+def _stable_uuid(text: str) -> str:
+    hex_str = hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]
+    return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:32]}"
+
+
+def _union_point_id(*, dataset_name: str, source_doc_id: str, union_namespace: str) -> str:
+    return _stable_uuid(f"{union_namespace}::{dataset_name}::{source_doc_id}")
+
+
+def _get_qdrant_env() -> Tuple[str, Optional[str]]:
+    url = os.getenv("QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("SIGIR_QDRANT_URL")
+    if not url:
+        raise SystemExit("QDRANT_URL not set")
+    key = (
+        os.getenv("QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("SIGIR_QDRANT_KEY")
+    )
+    return str(url), (str(key) if key else None)
+
+
+@dataclass(frozen=True)
+class RunSpec:
+    name: str
+    model: str
+    collection: str
+    torch_dtype: str
+    output_dtype: str
+
+
+SPECS: List[RunSpec] = [
+    RunSpec(
+        name="colqwen2.5_fp16_collection",
+        model="vidore/colqwen2.5-v0.2",
+        collection="vidore_beir_v2_3ds__colqwen25_v0_2__nocrop__union__fp16",
+        torch_dtype="float16",
+        output_dtype="float16",
+    ),
+    RunSpec(
+        name="colpali1.3_collection",
+        model="vidore/colpali-v1.3",
+        collection="vidore_beir_v2_3ds__colpali_v1_3__nocrop__union",
+        torch_dtype="float16",
+        output_dtype="float16",
+    ),
+]
+
+
+def _parse_dtype(s: str):
+    if s == "float16":
+        import torch
+
+        return torch.float16
+    if s == "float32":
+        import torch
+
+        return torch.float32
+    if s == "bfloat16":
+        import torch
+
+        return torch.bfloat16
+    return None
+
+
+def _np_dtype(s: str):
+    return np.float16 if s == "float16" else np.float32
+
+
+def _build_remapped_qrels(
+    *, corpus, qrels, dataset_name: str, collection: str
+) -> Dict[str, Dict[str, int]]:
+    # corpus doc_id values are stable_uuid(source_doc_id)
+    id_map: Dict[str, str] = {}
+    for doc in corpus:
+        source_doc_id = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
+        id_map[str(doc.doc_id)] = _union_point_id(
+            dataset_name=str(dataset_name),
+            source_doc_id=str(source_doc_id),
+            union_namespace=str(collection),
+        )
+
+    remapped: Dict[str, Dict[str, int]] = {}
+    for qid, rels in (qrels or {}).items():
+        out: Dict[str, int] = {}
+        for did, score in (rels or {}).items():
+            mapped = id_map.get(str(did))
+            if mapped:
+                out[str(mapped)] = int(score)
+        if out:
+            remapped[str(qid)] = out
+    return remapped
+
+
+def _rank_stats_for_query(
+    *, ranking: List[str], qrels: Dict[str, int], top_k: int
+) -> Dict[str, Any]:
+    relset = {did for did, s in (qrels or {}).items() if int(s) > 0}
+    best_rank = None
+    for i, did in enumerate(ranking[:top_k]):
+        if str(did) in relset:
+            best_rank = i + 1
+            break
+    return {
+        "num_relevant": int(len(relset)),
+        "best_rank": int(best_rank) if best_rank is not None else None,
+        "hit@1": bool(best_rank == 1),
+        "hit@5": bool(best_rank is not None and best_rank <= 5),
+        "hit@10": bool(best_rank is not None and best_rank <= 10),
+        "hit@100": bool(best_rank is not None and best_rank <= 100),
+    }
+
+
+def _run_one(
+    *,
+    spec: RunSpec,
+    dataset_name: str,
+    corpus,
+    queries,
+    qrels,
+    top_k: int,
+    max_queries: int,
+    prefer_grpc: bool,
+    timeout: int,
+) -> Dict[str, Any]:
+    url, key = _get_qdrant_env()
+
+    remapped_qrels = _build_remapped_qrels(
+        corpus=corpus, qrels=qrels, dataset_name=dataset_name, collection=spec.collection
+    )
+    # Keep only queries with at least one positive relevant doc
+    kept = [
+        q for q in queries if any(v > 0 for v in remapped_qrels.get(str(q.query_id), {}).values())
+    ]
+    kept = kept[: int(max_queries)] if int(max_queries) > 0 else kept
+
+    flt = qm.Filter(
+        must=[qm.FieldCondition(key="dataset", match=qm.MatchValue(value=str(dataset_name)))]
+    )
+
+    embedder = VisualEmbedder(
+        model_name=str(spec.model),
+        torch_dtype=_parse_dtype(spec.torch_dtype),
+        output_dtype=_np_dtype(spec.output_dtype),
+    )
+    retriever = MultiVectorRetriever(
+        collection_name=str(spec.collection),
+        model_name=str(spec.model),
+        embedder=embedder,
+        qdrant_url=url,
+        qdrant_api_key=key,
+        prefer_grpc=bool(prefer_grpc),
+        request_timeout=int(timeout),
+    )
+
+    per_query: Dict[str, Any] = {}
+    hits1 = hits5 = hits10 = hits100 = 0
+    best_ranks: List[int] = []
+    for q in kept:
+        qid = str(q.query_id)
+        rels = remapped_qrels.get(qid, {})
+        res = retriever.search(q.text, top_k=int(top_k), mode="single_full", filter_obj=flt)
+        ranking = [str(r["id"]) for r in (res or [])]
+        st = _rank_stats_for_query(ranking=ranking, qrels=rels, top_k=int(top_k))
+        per_query[qid] = {
+            "text": str(q.text),
+            "stats": st,
+            "top10": ranking[:10],
+        }
+        hits1 += 1 if st["hit@1"] else 0
+        hits5 += 1 if st["hit@5"] else 0
+        hits10 += 1 if st["hit@10"] else 0
+        hits100 += 1 if st["hit@100"] else 0
+        if st["best_rank"] is not None:
+            best_ranks.append(int(st["best_rank"]))
+
+    n = max(len(kept), 1)
+    summary = {
+        "queries_eval": int(len(kept)),
+        "hit_rate@1": float(hits1 / n),
+        "hit_rate@5": float(hits5 / n),
+        "hit_rate@10": float(hits10 / n),
+        "hit_rate@100": float(hits100 / n),
+        "median_best_rank": float(np.median(best_ranks)) if best_ranks else None,
+        "mean_best_rank": float(np.mean(best_ranks)) if best_ranks else None,
+    }
+
+    # Best-effort release memory
+    try:
+        import torch
+
+        del retriever
+        del embedder
+        gc.collect()
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+        elif torch.backends.mps.is_available():
+            torch.mps.empty_cache()
+    except Exception:
+        pass
+
+    return {
+        "spec": spec.__dict__,
+        "summary": summary,
+        "per_query": per_query,
+    }
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", default="vidore/esg_reports_v2")
+    ap.add_argument("--top-k", type=int, default=100)
+    ap.add_argument("--max-queries", type=int, default=50)
+    ap.add_argument("--prefer-grpc", action="store_true", default=True)
+    ap.add_argument("--timeout", type=int, default=120)
+    ap.add_argument("--out", default="auto")
+    args = ap.parse_args()
+
+    dataset_name = str(args.dataset)
+    corpus, queries, qrels = load_vidore_beir_dataset(dataset_name)
+
+    out_dir = Path("results") / "model_compare"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    out_path = out_dir / (
+        args.out
+        if str(args.out) != "auto"
+        else f"compare__{dataset_name.replace('/', '_')}__top{int(args.top_k)}__q{int(args.max_queries)}.json"
+    )
+
+    out: Dict[str, Any] = {
+        "dataset": dataset_name,
+        "top_k": int(args.top_k),
+        "max_queries": int(args.max_queries),
+        "runs": {},
+    }
+    for spec in SPECS:
+        out["runs"][spec.name] = _run_one(
+            spec=spec,
+            dataset_name=dataset_name,
+            corpus=corpus,
+            queries=queries,
+            qrels=qrels,
+            top_k=int(args.top_k),
+            max_queries=int(args.max_queries),
+            prefer_grpc=bool(args.prefer_grpc),
+            timeout=int(args.timeout),
+        )
+
+    out_path.write_text(json.dumps(out, ensure_ascii=False, indent=2))
+    print(f"Wrote: {out_path}")
+    print(json.dumps({k: v["summary"] for k, v in out["runs"].items()}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

scripts/create_qdrant_payload_indexes.py

@@ -0,0 +1,103 @@
+import argparse
+import os
+from pathlib import Path
+from typing import Dict, List
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _infer_type(values: List[object]) -> str:
+    for v in values:
+        if isinstance(v, bool):
+            return "bool"
+    for v in values:
+        if isinstance(v, int) and not isinstance(v, bool):
+            return "integer"
+    for v in values:
+        if isinstance(v, float):
+            return "float"
+    return "keyword"
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", type=str, required=True)
+    parser.add_argument("--prefer-grpc", action="store_true")
+    parser.add_argument("--sample", type=int, default=200)
+    parser.add_argument(
+        "--only",
+        type=str,
+        default="",
+        help="Comma-separated list of payload fields to index (optional).",
+    )
+    args = parser.parse_args()
+
+    _maybe_load_dotenv()
+
+    qdrant_url = os.getenv("QDRANT_URL")
+    if not qdrant_url:
+        raise ValueError("QDRANT_URL not set")
+    qdrant_api_key = os.getenv("QDRANT_API_KEY")
+
+    from qdrant_client import QdrantClient
+
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        check_compatibility=False,
+        timeout=120,
+    )
+
+    points, _ = client.scroll(
+        collection_name=args.collection,
+        limit=int(args.sample),
+        with_payload=True,
+        with_vectors=False,
+    )
+    if not points:
+        raise ValueError(f"No points found in collection '{args.collection}'")
+
+    only = [s.strip() for s in args.only.split(",") if s.strip()]
+    only_set = set(only) if only else None
+
+    values_by_key: Dict[str, List[object]] = {}
+    for p in points:
+        payload = p.payload or {}
+        if not isinstance(payload, dict):
+            continue
+        for k, v in payload.items():
+            if only_set is not None and k not in only_set:
+                continue
+            if isinstance(v, dict) or isinstance(v, list):
+                continue
+            values_by_key.setdefault(k, []).append(v)
+
+    from visual_rag.indexing.qdrant_indexer import QdrantIndexer
+
+    indexer = QdrantIndexer(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        collection_name=args.collection,
+        prefer_grpc=args.prefer_grpc,
+    )
+
+    fields = [{"field": k, "type": _infer_type(vs)} for k, vs in sorted(values_by_key.items())]
+    if not fields:
+        raise ValueError("No indexable payload fields found (all were nested or empty?)")
+
+    indexer.create_payload_indexes(fields=fields)
+    print(
+        f"Created/ensured {len(fields)} payload indexes on '{args.collection}': {[f['field'] for f in fields]}"
+    )
+
+
+if __name__ == "__main__":
+    main()

scripts/debug_failed_docs.py

@@ -0,0 +1,179 @@
+import argparse
+import json
+from pathlib import Path
+
+
+def _ensure_pil(img):
+    from PIL import Image
+
+    if img is None:
+        return None
+    if isinstance(img, Image.Image):
+        return img.convert("RGB")
+    try:
+        return img.convert("RGB")
+    except Exception:
+        raise TypeError(f"Unsupported image type: {type(img)}")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, required=True)
+    parser.add_argument("--model", type=str, default="vidore/colSmol-500M")
+    parser.add_argument("--device", type=str, default=None)
+    parser.add_argument(
+        "--torch-dtype", type=str, default="float16", choices=["float16", "float32", "bfloat16"]
+    )
+    parser.add_argument(
+        "--processor-speed", type=str, default="fast", choices=["fast", "slow", "auto"]
+    )
+    parser.add_argument("--source-doc-ids", type=str, nargs="+", required=True)
+    parser.add_argument("--crop-empty", action="store_true", default=False)
+    parser.add_argument("--crop-empty-percentage-to-remove", type=float, default=0.99)
+    parser.add_argument("--crop-empty-remove-page-number", action="store_true", default=False)
+    parser.add_argument("--crop-empty-preserve-border-px", type=int, default=1)
+    parser.add_argument("--crop-empty-uniform-std-threshold", type=float, default=0.0)
+    parser.add_argument("--out-dir", type=str, default="results/paper_eval/debug_failed_docs")
+    args = parser.parse_args()
+
+    import numpy as np
+    import torch
+
+    from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+    from visual_rag.embedding.visual_embedder import VisualEmbedder
+    from visual_rag.preprocessing.crop_empty import CropEmptyConfig, crop_empty
+
+    torch_dtype = {"float16": torch.float16, "float32": torch.float32, "bfloat16": torch.bfloat16}[
+        args.torch_dtype
+    ]
+
+    out_dir = Path(args.out_dir)
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    corpus, _, _ = load_vidore_beir_dataset(args.dataset)
+    wanted = set(str(x) for x in args.source_doc_ids)
+    found = []
+    for d in corpus:
+        sid = str((d.payload or {}).get("source_doc_id") or "")
+        if sid in wanted:
+            found.append(d)
+    found_by_id = {str((d.payload or {}).get("source_doc_id") or ""): d for d in found}
+    missing = [x for x in args.source_doc_ids if str(x) not in found_by_id]
+    if missing:
+        raise SystemExit(f"Could not find source_doc_id(s) in corpus: {missing}")
+
+    embedder = VisualEmbedder(
+        model_name=str(args.model),
+        device=args.device,
+        torch_dtype=torch_dtype,
+        processor_speed=str(args.processor_speed),
+        batch_size=1,
+    )
+
+    report = {}
+    for sid in args.source_doc_ids:
+        d = found_by_id[str(sid)]
+        original_img = _ensure_pil(d.image)
+        original_path = (
+            out_dir / f"{args.dataset.replace('/', '__')}__source_doc_id={sid}__original.png"
+        )
+        original_img.save(original_path)
+
+        crop_meta = {
+            "applied": False,
+            "crop_box": None,
+            "original_width": int(original_img.width),
+            "original_height": int(original_img.height),
+            "cropped_width": int(original_img.width),
+            "cropped_height": int(original_img.height),
+        }
+        embed_img = original_img
+        if bool(args.crop_empty):
+            embed_img, crop_meta = crop_empty(
+                original_img,
+                config=CropEmptyConfig(
+                    percentage_to_remove=float(args.crop_empty_percentage_to_remove),
+                    remove_page_number=bool(args.crop_empty_remove_page_number),
+                    preserve_border_px=int(args.crop_empty_preserve_border_px),
+                    uniform_rowcol_std_threshold=float(args.crop_empty_uniform_std_threshold),
+                ),
+            )
+
+        cropped_path = (
+            out_dir / f"{args.dataset.replace('/', '__')}__source_doc_id={sid}__cropped.png"
+        )
+        _ensure_pil(embed_img).save(cropped_path)
+
+        embeddings, token_infos = embedder.embed_images(
+            [embed_img],
+            batch_size=1,
+            return_token_info=True,
+            show_progress=False,
+        )
+        emb = embeddings[0]
+        token_info = token_infos[0] or {}
+
+        emb_np = (
+            emb.cpu().float().numpy() if hasattr(emb, "cpu") else np.array(emb, dtype=np.float32)
+        )
+        visual_indices = token_info.get("visual_token_indices") or list(range(int(emb_np.shape[0])))
+        visual_embedding = emb_np[visual_indices].astype(np.float32)
+
+        tile_pooled = embedder.mean_pool_visual_embedding(
+            visual_embedding, token_info, target_vectors=32
+        )
+        experimental_pooled = embedder.experimental_pool_visual_embedding(
+            visual_embedding,
+            token_info,
+            target_vectors=32,
+            mean_pool=tile_pooled,
+        )
+
+        n_rows = token_info.get("n_rows")
+        n_cols = token_info.get("n_cols")
+        num_tiles_from_info = token_info.get("num_tiles")
+        num_visual_tokens = token_info.get("num_visual_tokens")
+        num_tiles_from_tokens = int(visual_embedding.shape[0]) // 64 + (
+            1 if int(visual_embedding.shape[0]) % 64 else 0
+        )
+
+        report[str(sid)] = {
+            "dataset": str(args.dataset),
+            "model": str(args.model),
+            "source_doc_id": str(sid),
+            "doc_id": str(getattr(d, "doc_id", "")),
+            "payload_source_doc_id": str((d.payload or {}).get("source_doc_id") or ""),
+            "original_image": {
+                "path": str(original_path),
+                "width": int(original_img.width),
+                "height": int(original_img.height),
+            },
+            "crop_meta": crop_meta,
+            "cropped_image": {
+                "path": str(cropped_path),
+                "width": int(_ensure_pil(embed_img).width),
+                "height": int(_ensure_pil(embed_img).height),
+            },
+            "processor": {
+                "n_rows": None if n_rows is None else int(n_rows),
+                "n_cols": None if n_cols is None else int(n_cols),
+                "num_tiles": None if num_tiles_from_info is None else int(num_tiles_from_info),
+                "num_visual_tokens": None if num_visual_tokens is None else int(num_visual_tokens),
+                "visual_token_indices_len": int(len(visual_indices)),
+                "num_tiles_from_visual_tokens_div64": int(num_tiles_from_tokens),
+            },
+            "embeddings": {
+                "full_embedding_shape": [int(x) for x in emb_np.shape],
+                "visual_embedding_shape": [int(x) for x in visual_embedding.shape],
+                "mean_pool_shape": [int(x) for x in tile_pooled.shape],
+                "experimental_pool_shape": [int(x) for x in experimental_pooled.shape],
+            },
+        }
+
+    out_json = out_dir / f"{args.dataset.replace('/', '__')}__debug_report.json"
+    out_json.write_text(json.dumps(report, indent=2), encoding="utf-8")
+    print(str(out_json))
+
+
+if __name__ == "__main__":
+    main()

scripts/debug_vidore_qrels_alignment.py

@@ -0,0 +1,189 @@
+"""
+Debug ViDoRe-v2 evaluation mismatches between qrels and Qdrant point IDs.
+
+This script helps answer:
+- Are relevant docs (from qrels) actually present in Qdrant?
+- Are we mapping qrels doc IDs to the correct Qdrant point IDs?
+- Does per_dataset filtering actually reduce the search space?
+- If docs exist, at what rank do they appear for single_full retrieval?
+
+Typical use:
+  python scripts/debug_vidore_qrels_alignment.py \\
+    --dataset vidore/esg_reports_v2 \\
+    --collection vidore_beir_v2_3ds__colqwen25_v0_2__nocrop__union__fp32 \\
+    --model vidore/colqwen2.5-v0.2 \\
+    --max-queries 5 \\
+    --top-k 200 \\
+    --no-prefer-grpc
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Dict, Tuple
+
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qm
+
+from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+from visual_rag import VisualEmbedder
+from visual_rag.retrieval import MultiVectorRetriever
+
+
+def _stable_uuid(text: str) -> str:
+    import hashlib
+
+    hex_str = hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]
+    return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:32]}"
+
+
+def _union_point_id(*, dataset_name: str, source_doc_id: str, union_namespace: str) -> str:
+    return _stable_uuid(f"{union_namespace}::{dataset_name}::{source_doc_id}")
+
+
+def _infer_qdrant_conn(prefer_grpc: bool, timeout: int) -> QdrantClient:
+    url = os.getenv("QDRANT_URL") or os.getenv("DEST_QDRANT_URL") or os.getenv("SIGIR_QDRANT_URL")
+    if not url:
+        raise SystemExit("QDRANT_URL not set")
+    key = (
+        os.getenv("QDRANT_API_KEY")
+        or os.getenv("DEST_QDRANT_API_KEY")
+        or os.getenv("SIGIR_QDRANT_KEY")
+    )
+    return QdrantClient(
+        url=url,
+        api_key=key,
+        prefer_grpc=bool(prefer_grpc),
+        timeout=int(timeout),
+        check_compatibility=False,
+    )
+
+
+def _count_by_dataset(client: QdrantClient, collection: str, dataset: str) -> Tuple[int, int]:
+    # exact counts can be slow; we keep it exact for correctness.
+    all_cnt = client.count(collection_name=collection, exact=True).count
+    ds_cnt = client.count(
+        collection_name=collection,
+        count_filter=qm.Filter(
+            must=[qm.FieldCondition(key="dataset", match=qm.MatchValue(value=str(dataset)))]
+        ),
+        exact=True,
+    ).count
+    return int(ds_cnt), int(all_cnt)
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dataset", required=True)
+    ap.add_argument("--collection", required=True)
+    ap.add_argument("--model", required=True)
+    ap.add_argument("--top-k", type=int, default=200)
+    ap.add_argument("--max-queries", type=int, default=5)
+    ap.add_argument("--prefer-grpc", action="store_true", default=False)
+    ap.add_argument("--timeout", type=int, default=120)
+    ap.add_argument(
+        "--torch-dtype", default="auto", choices=["auto", "float32", "float16", "bfloat16"]
+    )
+    args = ap.parse_args()
+
+    corpus, queries, qrels = load_vidore_beir_dataset(str(args.dataset))
+    print(
+        f"Loaded dataset={args.dataset}: corpus={len(corpus)} queries={len(queries)} qrels_q={len(qrels)}"
+    )
+
+    # Build mapping exactly like the benchmark does.
+    id_map: Dict[str, str] = {}
+    for doc in corpus:
+        src = str((doc.payload or {}).get("source_doc_id") or doc.doc_id)
+        id_map[str(doc.doc_id)] = _union_point_id(
+            dataset_name=str(args.dataset),
+            source_doc_id=str(src),
+            union_namespace=str(args.collection),
+        )
+
+    remapped_qrels: Dict[str, Dict[str, int]] = {}
+    for qid, rels in qrels.items():
+        out: Dict[str, int] = {}
+        for did, score in rels.items():
+            if int(score) <= 0:
+                continue
+            mapped = id_map.get(str(did))
+            if mapped:
+                out[str(mapped)] = int(score)
+        if out:
+            remapped_qrels[str(qid)] = out
+
+    # Connectivity + counts
+    client = _infer_qdrant_conn(bool(args.prefer_grpc), int(args.timeout))
+    ds_cnt, all_cnt = _count_by_dataset(client, str(args.collection), str(args.dataset))
+    print(f"Qdrant counts: dataset={ds_cnt} / all={all_cnt} (collection={args.collection})")
+
+    # Pick queries that still have qrels after remap
+    kept = [q for q in queries if str(q.query_id) in remapped_qrels]
+    kept = kept[: int(args.max_queries)]
+    print(f"Queries kept after qrels remap: {len(kept)} (showing up to {args.max_queries})")
+
+    # Build retriever with the exact same embedder/retrieval path.
+    td = None
+    if str(args.torch_dtype) != "auto":
+        import torch
+
+        td = {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[
+            str(args.torch_dtype)
+        ]
+    embedder = VisualEmbedder(model_name=str(args.model), torch_dtype=td)
+    retriever = MultiVectorRetriever(
+        collection_name=str(args.collection),
+        model_name=str(args.model),
+        embedder=embedder,
+        qdrant_client=client,
+        prefer_grpc=bool(args.prefer_grpc),
+        request_timeout=int(args.timeout),
+    )
+
+    flt = qm.Filter(
+        must=[qm.FieldCondition(key="dataset", match=qm.MatchValue(value=str(args.dataset)))]
+    )
+
+    for i, q in enumerate(kept):
+        qid = str(q.query_id)
+        rels = remapped_qrels.get(qid, {})
+        # Only positive qrels are truly relevant.
+        rel_ids = [rid for rid, s in (rels or {}).items() if int(s) > 0]
+        print("\n" + "-" * 90)
+        print(f"Q{i}: {qid}  text={q.text[:120]!r}")
+        print(f"  relevant_ids(remapped)={len(rel_ids)}  sample={rel_ids[:3]}")
+
+        # Check if relevant IDs exist in Qdrant at all
+        exists = 0
+        try:
+            recs = client.retrieve(
+                collection_name=str(args.collection),
+                ids=rel_ids[:20],
+                with_payload=False,
+                with_vectors=False,
+                timeout=int(args.timeout),
+            )
+            exists = len(recs)
+        except Exception:
+            exists = 0
+        print(f"  relevant_ids_exist_in_qdrant(sample<=20): {exists}")
+
+        # Search per_dataset filter
+        res = retriever.search(q.text, top_k=int(args.top_k), mode="single_full", filter_obj=flt)
+        ranked = [str(r["id"]) for r in res]
+        # Find best rank of any relevant doc
+        best_rank = None
+        for rid in rel_ids:
+            if rid in ranked:
+                rnk = ranked.index(rid) + 1
+                best_rank = rnk if best_rank is None else min(best_rank, rnk)
+        print(f"  best_rank_in_top{args.top_k} (per_dataset filter): {best_rank}")
+        print(f"  top10 ids: {ranked[:10]}")
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

scripts/dedupe_failure_logs.py

@@ -0,0 +1,70 @@
+import argparse
+import json
+import os
+import tempfile
+from pathlib import Path
+
+
+def _iter_jsonl(path: Path):
+    if not path.exists():
+        return
+    with path.open("r", encoding="utf-8") as f:
+        for line in f:
+            s = (line or "").strip()
+            if not s:
+                continue
+            yield json.loads(s)
+
+
+def _key(obj: dict) -> str:
+    for k in ("union_doc_id", "id", "doc_id", "source_doc_id"):
+        v = obj.get(k)
+        if v:
+            return str(v)
+    return json.dumps(obj, sort_keys=True)
+
+
+def _write_atomic(path: Path, lines: list[str]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    fd, tmp = tempfile.mkstemp(prefix=path.name + ".", dir=str(path.parent))
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as f:
+            for ln in lines:
+                f.write(ln)
+                f.write("\n")
+        os.replace(tmp, path)
+    finally:
+        try:
+            if os.path.exists(tmp):
+                os.unlink(tmp)
+        except Exception:
+            pass
+
+
+def dedupe_jsonl(path: Path) -> dict:
+    last_by_key: dict[str, dict] = {}
+    order: list[str] = []
+    for obj in _iter_jsonl(path):
+        k = _key(obj)
+        if k not in last_by_key:
+            order.append(k)
+        last_by_key[k] = obj
+
+    out_lines = [json.dumps(last_by_key[k], ensure_ascii=False) for k in order]
+    _write_atomic(path, out_lines)
+    return {"path": str(path), "unique": len(out_lines)}
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--paths", type=str, nargs="+", required=True)
+    args = parser.parse_args()
+
+    for p in args.paths:
+        path = Path(p)
+        res = dedupe_jsonl(path)
+        print(f"{res['path']}: unique={res['unique']}")
+
+
+if __name__ == "__main__":
+    main()

scripts/force_qdrant_reindex.py

@@ -0,0 +1,119 @@
+import argparse
+import os
+import time
+from pathlib import Path
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _indexed_total(indexed_vectors_count) -> int:
+    if indexed_vectors_count is None:
+        return 0
+    if isinstance(indexed_vectors_count, dict):
+        try:
+            return int(sum(int(v) for v in indexed_vectors_count.values()))
+        except Exception:
+            return 0
+    try:
+        return int(indexed_vectors_count)
+    except Exception:
+        return 0
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", type=str, required=True)
+    parser.add_argument("--prefer-grpc", action="store_true")
+    parser.add_argument("--url", type=str, default="")
+    parser.add_argument("--api-key", type=str, default="")
+    parser.add_argument("--indexing-threshold", type=int, default=0)
+    parser.add_argument("--m", type=int, default=32)
+    parser.add_argument("--ef-construct", type=int, default=100)
+    parser.add_argument("--full-scan-threshold", type=int, default=10000)
+    parser.add_argument(
+        "--on-disk",
+        action="store_true",
+        help="Store HNSW index on disk (recommended for large vectors).",
+    )
+    parser.add_argument("--max-indexing-threads", type=int, default=0)
+    parser.add_argument("--wait", action="store_true")
+    parser.add_argument("--timeout-sec", type=int, default=600)
+    parser.add_argument("--poll-sec", type=float, default=2.0)
+    args = parser.parse_args()
+
+    _maybe_load_dotenv()
+
+    qdrant_url = args.url or os.getenv("QDRANT_URL")
+    if not qdrant_url:
+        raise ValueError("QDRANT_URL not set")
+    qdrant_api_key = args.api_key or os.getenv("QDRANT_API_KEY")
+
+    from qdrant_client import QdrantClient
+    from qdrant_client.http import models
+
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        check_compatibility=False,
+        timeout=120,
+    )
+
+    hnsw = models.HnswConfigDiff(
+        m=int(args.m),
+        ef_construct=int(args.ef_construct),
+        full_scan_threshold=int(args.full_scan_threshold),
+        on_disk=bool(args.on_disk),
+        max_indexing_threads=int(args.max_indexing_threads),
+    )
+
+    vectors_config = {
+        "initial": models.VectorParamsDiff(hnsw_config=hnsw, on_disk=True),
+        "mean_pooling": models.VectorParamsDiff(hnsw_config=hnsw),
+        "global_pooling": models.VectorParamsDiff(hnsw_config=hnsw),
+    }
+
+    client.update_collection(
+        collection_name=args.collection,
+        optimizers_config=models.OptimizersConfigDiff(
+            indexing_threshold=int(args.indexing_threshold)
+        ),
+        hnsw_config=hnsw,
+        vectors_config=vectors_config,
+    )
+
+    info = client.get_collection(args.collection)
+    print(
+        f"Triggered reindex update for '{args.collection}'. "
+        f"points={info.points_count}, indexed_vectors={info.indexed_vectors_count}, "
+        f"status={getattr(getattr(info.status, 'value', None), 'value', getattr(info, 'status', None))}"
+    )
+
+    if not args.wait:
+        return
+
+    start = time.time()
+    while True:
+        info = client.get_collection(args.collection)
+        indexed_total = _indexed_total(info.indexed_vectors_count)
+        total = int(info.points_count or 0)
+        print(
+            f"poll: points={info.points_count}, indexed_vectors={info.indexed_vectors_count}, "
+            f"segments={getattr(info, 'segments_count', None)}"
+        )
+        if total > 0 and indexed_total >= total:
+            return
+        if time.time() - start > args.timeout_sec:
+            return
+        time.sleep(max(0.1, float(args.poll_sec)))
+
+
+if __name__ == "__main__":
+    main()

scripts/inspect_qdrant_collection.py

@@ -0,0 +1,94 @@
+import argparse
+import json
+import os
+from pathlib import Path
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _as_jsonable(obj):
+    if obj is None:
+        return None
+    if isinstance(obj, (str, int, float, bool)):
+        return obj
+    if isinstance(obj, dict):
+        return {str(k): _as_jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_as_jsonable(v) for v in obj]
+    if hasattr(obj, "model_dump"):
+        try:
+            return obj.model_dump()
+        except Exception:
+            pass
+    if hasattr(obj, "__dict__"):
+        try:
+            return {k: _as_jsonable(v) for k, v in obj.__dict__.items()}
+        except Exception:
+            pass
+    return str(obj)
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", type=str, required=True)
+    parser.add_argument("--prefer-grpc", action="store_true")
+    parser.add_argument("--url", type=str, default="")
+    parser.add_argument("--api-key", type=str, default="")
+    parser.add_argument("--out", type=str, default="")
+    args = parser.parse_args()
+
+    _maybe_load_dotenv()
+
+    qdrant_url = args.url or os.getenv("QDRANT_URL")
+    if not qdrant_url:
+        raise ValueError("QDRANT_URL not set")
+    qdrant_api_key = args.api_key or os.getenv("QDRANT_API_KEY")
+
+    from qdrant_client import QdrantClient
+
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        check_compatibility=False,
+        timeout=120,
+    )
+
+    info = client.get_collection(args.collection)
+    payload_schema = getattr(info, "payload_schema", None)
+    snap = {
+        "collection": args.collection,
+        "points_count": _as_jsonable(getattr(info, "points_count", None)),
+        "indexed_vectors_count": _as_jsonable(getattr(info, "indexed_vectors_count", None)),
+        "segments_count": _as_jsonable(getattr(info, "segments_count", None)),
+        "status": _as_jsonable(
+            getattr(getattr(info, "status", None), "value", getattr(info, "status", None))
+        ),
+        "optimizer_status": _as_jsonable(
+            getattr(
+                getattr(info, "optimizer_status", None),
+                "value",
+                getattr(info, "optimizer_status", None),
+            )
+        ),
+        "config": _as_jsonable(getattr(info, "config", None)),
+        "payload_schema": _as_jsonable(payload_schema),
+    }
+
+    print(json.dumps(snap, indent=2)[:10000])
+    if args.out:
+        out_path = Path(args.out)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump(snap, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_clone_collection_no_index.py

@@ -0,0 +1,253 @@
+"""
+Clone an existing Qdrant collection into a new collection with indexing disabled.
+
+Why: Qdrant doesn't provide an in-place "de-index" for already-built HNSW.
+This script clones points (payload + vectors) into a fresh collection created
+with a very large `indexing_threshold`, so `indexed_vectors_count` stays 0.
+
+Usage:
+  python scripts/qdrant_clone_collection_no_index.py \
+    --source vidore_beir_v2_... \
+    --dest vidore_beir_v2_...__noindex \
+    --embedding-dim 128 \
+    --vector-dtype float32 \
+    --indexing-threshold 1000000000 \
+    --prefer-grpc
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from typing import Any, Dict, List, Optional, Sequence
+
+try:
+    from dotenv import load_dotenv
+
+    DOTENV_AVAILABLE = True
+except Exception:
+    DOTENV_AVAILABLE = False
+
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qm
+
+
+def _get_env(name: str) -> Optional[str]:
+    v = os.getenv(name)
+    if v is None:
+        return None
+    v = str(v).strip()
+    return v or None
+
+
+def _require(value: Optional[str], *, name: str) -> str:
+    if not value:
+        raise SystemExit(
+            f"Missing {name}. Set it in env (preferred) or pass the corresponding flag."
+        )
+    return value
+
+
+def _chunks(seq: Sequence[Any], n: int) -> List[Sequence[Any]]:
+    return [seq[i : i + n] for i in range(0, len(seq), n)]
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", required=True, help="Existing source collection name")
+    parser.add_argument("--dest", required=True, help="Destination collection name")
+    parser.add_argument(
+        "--qdrant-url",
+        default=None,
+        help="Qdrant URL (or set QDRANT_URL env var)",
+    )
+    parser.add_argument(
+        "--qdrant-api-key",
+        default=None,
+        help="Qdrant API key (or set QDRANT_API_KEY env var)",
+    )
+    parser.add_argument(
+        "--prefer-grpc",
+        action="store_true",
+        help="Use gRPC transport (recommended for large vectors)",
+    )
+    parser.add_argument(
+        "--timeout",
+        type=float,
+        default=300.0,
+        help="Request timeout seconds",
+    )
+    parser.add_argument(
+        "--embedding-dim",
+        type=int,
+        default=128,
+        help="Vector dimension (typically 128 for ColPali/ColQwen)",
+    )
+    parser.add_argument(
+        "--vector-dtype",
+        choices=["float16", "float32"],
+        default="float32",
+        help="Vector datatype for destination collection",
+    )
+    parser.add_argument(
+        "--indexing-threshold",
+        type=int,
+        default=1_000_000_000,
+        help="Large value prevents HNSW building for small collections",
+    )
+    # Note: some qdrant-client versions don't support full_scan_threshold in OptimizersConfigDiff.
+    parser.add_argument(
+        "--recreate-dest",
+        action="store_true",
+        help="Delete destination collection if it already exists",
+    )
+    parser.add_argument(
+        "--scroll-limit",
+        type=int,
+        default=256,
+        help="How many points to fetch per scroll call",
+    )
+    parser.add_argument(
+        "--upsert-batch-size",
+        type=int,
+        default=64,
+        help="How many points per upsert call",
+    )
+
+    args = parser.parse_args()
+
+    if DOTENV_AVAILABLE:
+        load_dotenv()
+
+    url = args.qdrant_url or _get_env("QDRANT_URL")
+    api_key = args.qdrant_api_key or _get_env("QDRANT_API_KEY")
+    url = _require(url, name="QDRANT_URL/--qdrant-url")
+    api_key = _require(api_key, name="QDRANT_API_KEY/--qdrant-api-key")
+
+    client = QdrantClient(
+        url=url,
+        api_key=api_key,
+        prefer_grpc=bool(args.prefer_grpc),
+        timeout=float(args.timeout),
+    )
+
+    # Verify source exists
+    src_info = client.get_collection(args.source)
+    print(f"✅ Source collection found: {args.source}")
+    print(
+        f"   points_count≈{src_info.points_count} indexed_vectors_count={src_info.indexed_vectors_count}"
+    )
+
+    # Create/recreate destination with the same named vectors layout
+    if args.recreate_dest:
+        try:
+            client.delete_collection(args.dest)
+            print(f"🗑️ Deleted existing destination: {args.dest}")
+        except Exception as e:
+            print(f"⚠️ Could not delete dest (may not exist): {e}")
+
+    # Use same vector names as toolkit expects
+    datatype = qm.Datatype.FLOAT16 if args.vector_dtype == "float16" else qm.Datatype.FLOAT32
+    multivector_config = qm.MultiVectorConfig(comparator=qm.MultiVectorComparator.MAX_SIM)
+    vectors_config: Dict[str, qm.VectorParams] = {
+        "initial": qm.VectorParams(
+            size=int(args.embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=True,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "mean_pooling": qm.VectorParams(
+            size=int(args.embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "experimental_pooling": qm.VectorParams(
+            size=int(args.embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "global_pooling": qm.VectorParams(
+            size=int(args.embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            datatype=datatype,
+        ),
+    }
+
+    try:
+        client.create_collection(
+            collection_name=args.dest,
+            vectors_config=vectors_config,
+            optimizers_config=qm.OptimizersConfigDiff(
+                indexing_threshold=int(args.indexing_threshold),
+            ),
+        )
+        # Keep filename index for skip_existing (cheap and helpful)
+        try:
+            client.create_payload_index(
+                collection_name=args.dest,
+                field_name="filename",
+                field_schema=qm.PayloadSchemaType.KEYWORD,
+            )
+        except Exception:
+            pass
+        print(f"✅ Created destination collection: {args.dest}")
+    except Exception as e:
+        # If it already exists, proceed (unless user expected recreate)
+        print(f"ℹ️ Destination create skipped/failed (may already exist): {e}")
+
+    # Clone points
+    next_offset = None
+    total = 0
+
+    while True:
+        points, next_offset = client.scroll(
+            collection_name=args.source,
+            limit=int(args.scroll_limit),
+            with_payload=True,
+            with_vectors=True,
+            offset=next_offset,
+        )
+        if not points:
+            break
+
+        # Upsert in batches
+        for batch in _chunks(points, int(args.upsert_batch_size)):
+            upsert_points: List[qm.PointStruct] = []
+            for p in batch:
+                # p.vector may be dict (named vectors) or list (single). We expect dict.
+                vectors = p.vector
+                payload = p.payload or {}
+                upsert_points.append(
+                    qm.PointStruct(
+                        id=p.id,
+                        vector=vectors,
+                        payload=payload,
+                    )
+                )
+
+            client.upsert(
+                collection_name=args.dest,
+                points=upsert_points,
+                wait=True,
+            )
+            total += len(upsert_points)
+            if total % 500 == 0:
+                print(f"… cloned {total} points")
+
+    dst_info = client.get_collection(args.dest)
+    exact = client.count(collection_name=args.dest, exact=True)
+    print("✅ Clone complete")
+    print(
+        f"   dest.points_count≈{dst_info.points_count} dest.indexed_vectors_count={dst_info.indexed_vectors_count}"
+    )
+    print(f"   dest.count(exact)= {exact.count}")
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_debug_collection.py

@@ -0,0 +1,257 @@
+"""
+Qdrant collection debugging / inspection helpers.
+
+This script is intentionally lightweight and safe:
+- Read-only operations (count/scroll/retrieve)
+- Prints exact vs approximate counts (Qdrant UI often shows approximate)
+- Can verify that IDs listed in index_failures__*.jsonl logs are actually present
+
+Examples:
+
+  # Inspect counts + vector sanity (REST)
+  python scripts/qdrant_debug_collection.py \\
+    --collection vidore_beir_v2_3ds__colqwen25_v0_2__nocrop__union__fp32__grpc
+
+  # Same, but via gRPC
+  python scripts/qdrant_debug_collection.py \\
+    --collection vidore_beir_v2_3ds__colqwen25_v0_2__nocrop__union__fp32__grpc \\
+    --prefer-grpc
+
+  # Count per dataset (exact)
+  python scripts/qdrant_debug_collection.py \\
+    --collection <COLLECTION> \\
+    --datasets vidore/esg_reports_v2 vidore/biomedical_lectures_v2 vidore/economics_reports_v2
+
+  # Verify that any IDs in index_failures logs are present in Qdrant
+  python scripts/qdrant_debug_collection.py \\
+    --collection <COLLECTION> \\
+    --check-failures
+
+Environment:
+  export QDRANT_URL=...
+  export QDRANT_API_KEY=...  # optional
+
+Or create a .env in repo root with the same variables.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from collections import Counter
+from pathlib import Path
+from typing import Iterable
+
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qm
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except Exception:
+        return
+
+    for p in (Path(".env"), Path("..") / ".env"):
+        if p.exists():
+            load_dotenv(p)
+
+
+def _chunks(xs: list[str], n: int) -> Iterable[list[str]]:
+    for i in range(0, len(xs), n):
+        yield xs[i : i + n]
+
+
+def inspect_collection(*, client: QdrantClient, collection: str, sample_points: int) -> None:
+    info = client.get_collection(collection)
+    print("collection:", collection)
+    print("status:", info.status)
+    print("optimizer_status:", info.optimizer_status)
+    print("indexed_vectors_count:", info.indexed_vectors_count)
+    print()
+
+    approx = client.count(collection_name=collection, exact=False).count
+    exact = client.count(collection_name=collection, exact=True).count
+    print("info.points_count (approx):", info.points_count)
+    print("count(exact=False):", approx)
+    print("count(exact=True): ", exact)
+
+    if sample_points > 0:
+        points, _ = client.scroll(
+            collection_name=collection,
+            limit=int(sample_points),
+            with_payload=True,
+            with_vectors=True,
+        )
+
+        print("\nSample points vector sanity:")
+        for p in points:
+            vecs = p.vector or {}
+
+            def _len(v):
+                try:
+                    return len(v)
+                except Exception:
+                    return None
+
+            lengths = {
+                k: _len(v)
+                for k, v in vecs.items()
+                if k in {"initial", "mean_pooling", "experimental_pooling", "global_pooling"}
+            }
+            print("id:", p.id)
+            print("  dataset:", (p.payload or {}).get("dataset"))
+            print("  vector_keys:", sorted(list(vecs.keys())))
+            print("  lengths:", lengths)
+
+
+def count_per_dataset(*, client: QdrantClient, collection: str, datasets: list[str]) -> None:
+    if not datasets:
+        return
+    print("\nper-dataset exact counts:")
+    total = 0
+    for ds in datasets:
+        c = client.count(
+            collection_name=collection,
+            count_filter=qm.Filter(
+                must=[
+                    qm.FieldCondition(key="dataset", match=qm.MatchValue(value=str(ds))),
+                ]
+            ),
+            exact=True,
+        ).count
+        print(f"- {ds}: {c}")
+        total += int(c)
+    print("sum_datasets_exact:", total)
+
+
+def dataset_distribution_scroll(*, client: QdrantClient, collection: str, limit: int) -> None:
+    values: Counter[str] = Counter()
+    offset = None
+    seen = 0
+    while True:
+        points, next_offset = client.scroll(
+            collection_name=collection,
+            limit=min(int(limit), 2048),
+            offset=offset,
+            with_payload=True,
+            with_vectors=False,
+        )
+        if not points:
+            break
+        for p in points:
+            ds = (p.payload or {}).get("dataset")
+            values[str(ds)] += 1
+        seen += len(points)
+        offset = next_offset
+        if next_offset is None or (limit and seen >= int(limit)):
+            break
+
+    print("\nscroll distribution (dataset field):")
+    print("scrolled_points:", seen)
+    for k, v in values.most_common(20):
+        print(" ", k, v)
+
+
+def check_failure_logs_present(
+    *,
+    client: QdrantClient,
+    collection: str,
+    results_dir: Path,
+    retrieve_batch: int,
+) -> None:
+    base = results_dir / collection
+    if not base.exists():
+        raise SystemExit(f"results dir not found: {base}")
+
+    log_paths = sorted(base.glob("index_failures__*.jsonl"))
+    if not log_paths:
+        print("\nNo failure logs found under:", base)
+        return
+
+    failed_ids: set[str] = set()
+    for p in log_paths:
+        for line in p.read_text().splitlines():
+            line = (line or "").strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+            except Exception:
+                continue
+            u = obj.get("union_doc_id")
+            if u:
+                failed_ids.add(str(u))
+
+    print("\nfailure logs:")
+    for p in log_paths:
+        print(" -", p)
+    print("failed_ids_in_logs:", len(failed_ids))
+
+    missing: list[str] = []
+    ids = list(failed_ids)
+    for chunk in _chunks(ids, int(retrieve_batch)):
+        pts = client.retrieve(
+            collection_name=collection,
+            ids=chunk,
+            with_payload=False,
+            with_vectors=False,
+        )
+        present = set(str(p.id) for p in pts)
+        for pid in chunk:
+            if pid not in present:
+                missing.append(pid)
+
+    print("failed_ids_missing_in_qdrant:", len(missing))
+    if missing:
+        print("sample_missing_ids:", missing[:10])
+
+
+def main() -> None:
+    p = argparse.ArgumentParser(description="Qdrant collection debug utilities")
+    p.add_argument("--collection", required=True)
+    p.add_argument("--prefer-grpc", action="store_true", default=False)
+    p.add_argument("--timeout", type=int, default=120)
+    p.add_argument("--datasets", nargs="*", default=[])
+    p.add_argument("--sample-points", type=int, default=5)
+    p.add_argument("--scroll-limit", type=int, default=0, help="0 = no full scroll distribution")
+    p.add_argument("--check-failures", action="store_true", default=False)
+    p.add_argument("--results-dir", type=str, default="results")
+    p.add_argument("--retrieve-batch", type=int, default=64)
+    args = p.parse_args()
+
+    _maybe_load_dotenv()
+    url = os.getenv("QDRANT_URL")
+    key = os.getenv("QDRANT_API_KEY")
+    if not url:
+        raise SystemExit("QDRANT_URL not set")
+
+    client = QdrantClient(
+        url=url,
+        api_key=key,
+        prefer_grpc=bool(args.prefer_grpc),
+        timeout=int(args.timeout),
+    )
+
+    inspect_collection(
+        client=client, collection=str(args.collection), sample_points=int(args.sample_points)
+    )
+    count_per_dataset(client=client, collection=str(args.collection), datasets=list(args.datasets))
+
+    if int(args.scroll_limit) > 0:
+        dataset_distribution_scroll(
+            client=client, collection=str(args.collection), limit=int(args.scroll_limit)
+        )
+
+    if bool(args.check_failures):
+        check_failure_logs_present(
+            client=client,
+            collection=str(args.collection),
+            results_dir=Path(str(args.results_dir)),
+            retrieve_batch=int(args.retrieve_batch),
+        )
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_disable_hnsw.py

@@ -0,0 +1,192 @@
+"""
+Disable dense (HNSW) indexing for a Qdrant collection (make indexed_vectors_count go to 0).
+
+Qdrant supports disabling HNSW by setting `hnsw_config.m = 0`.
+For an already-indexed collection, Qdrant may run a reconstruction/optimization pass.
+Once that finishes, `indexed_vectors_count` should become 0.
+
+Ref: "Optimizing Memory for Bulk Uploads" (Qdrant, Feb 2025) recommends `m=0` to disable HNSW.
+
+Usage:
+  python scripts/qdrant_disable_hnsw.py --collection "my_collection" --wait
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+from typing import Any, Optional
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except Exception:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _get_env(name: str) -> Optional[str]:
+    v = os.getenv(name)
+    if v is None:
+        return None
+    v = str(v).strip()
+    return v or None
+
+
+def _as_jsonable(obj: Any):
+    if obj is None:
+        return None
+    if isinstance(obj, (str, int, float, bool)):
+        return obj
+    if isinstance(obj, dict):
+        return {str(k): _as_jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_as_jsonable(v) for v in obj]
+    if hasattr(obj, "model_dump"):
+        try:
+            return obj.model_dump()
+        except Exception:
+            pass
+    return str(obj)
+
+
+def _indexed_total(indexed_vectors_count) -> int:
+    if indexed_vectors_count is None:
+        return 0
+    if isinstance(indexed_vectors_count, dict):
+        try:
+            return int(sum(int(v) for v in indexed_vectors_count.values()))
+        except Exception:
+            return 0
+    try:
+        return int(indexed_vectors_count)
+    except Exception:
+        return 0
+
+
+def _snapshot(client, collection: str) -> dict:
+    info = client.get_collection(collection)
+    status = getattr(info, "status", None)
+    if hasattr(status, "value"):
+        status = status.value
+    optimizer_status = getattr(info, "optimizer_status", None)
+    if hasattr(optimizer_status, "value"):
+        optimizer_status = optimizer_status.value
+    return {
+        "status": _as_jsonable(status),
+        "optimizer_status": _as_jsonable(optimizer_status),
+        "points_count": _as_jsonable(getattr(info, "points_count", None)),
+        "indexed_vectors_count": _as_jsonable(getattr(info, "indexed_vectors_count", None)),
+        "segments_count": _as_jsonable(getattr(info, "segments_count", None)),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", required=True)
+    parser.add_argument("--prefer-grpc", action="store_true")
+    parser.add_argument("--url", default="")
+    parser.add_argument("--api-key", default="")
+    parser.add_argument("--timeout", type=float, default=120.0)
+    parser.add_argument("--wait", action="store_true")
+    parser.add_argument("--poll-sec", type=float, default=5.0)
+    parser.add_argument("--timeout-sec", type=float, default=1800.0)
+    parser.add_argument("--dump-json", default="", help="Optional path to dump snapshots JSON")
+    args = parser.parse_args()
+
+    _maybe_load_dotenv()
+
+    qdrant_url = args.url or _get_env("QDRANT_URL")
+    if not qdrant_url:
+        raise SystemExit("QDRANT_URL not set (or pass --url)")
+    qdrant_api_key = args.api_key or _get_env("QDRANT_API_KEY")
+
+    from qdrant_client import QdrantClient
+    from qdrant_client.http import models
+
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=bool(args.prefer_grpc),
+        check_compatibility=False,
+        timeout=float(args.timeout),
+    )
+
+    before = _snapshot(client, args.collection)
+    print(
+        f"Before: points={before['points_count']} indexed_vectors={before['indexed_vectors_count']} "
+        f"status={before['status']} optimizer={before['optimizer_status']} segments={before['segments_count']}"
+    )
+
+    # Disable HNSW for dense vectors
+    client.update_collection(
+        collection_name=args.collection,
+        hnsw_config=models.HnswConfigDiff(m=0),
+    )
+    after = _snapshot(client, args.collection)
+    print(
+        f"After update(m=0): points={after['points_count']} indexed_vectors={after['indexed_vectors_count']} "
+        f"status={after['status']} optimizer={after['optimizer_status']} segments={after['segments_count']}"
+    )
+
+    if args.dump_json:
+        out_path = Path(args.dump_json)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump(
+                {
+                    "collection": args.collection,
+                    "before": before,
+                    "after_update": after,
+                },
+                f,
+                indent=2,
+            )
+
+    if not args.wait:
+        return
+
+    start = time.time()
+    while True:
+        snap = _snapshot(client, args.collection)
+        indexed = _indexed_total(snap["indexed_vectors_count"])
+        if indexed == 0:
+            print(
+                f"✅ Done: indexed_vectors_count is 0. points={snap['points_count']} "
+                f"status={snap['status']} optimizer={snap['optimizer_status']}"
+            )
+            if args.dump_json:
+                out_path = Path(args.dump_json)
+                out_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(out_path, "w") as f:
+                    json.dump(
+                        {
+                            "collection": args.collection,
+                            "before": before,
+                            "after_update": after,
+                            "complete": snap,
+                        },
+                        f,
+                        indent=2,
+                    )
+            return
+        if time.time() - start > float(args.timeout_sec):
+            print(
+                f"⏳ Timeout waiting for indexed_vectors_count=0. indexed_vectors={snap['indexed_vectors_count']}, "
+                f"points={snap['points_count']}, status={snap['status']}, optimizer={snap['optimizer_status']}"
+            )
+            return
+        print(
+            f"… waiting: indexed_vectors={snap['indexed_vectors_count']} points={snap['points_count']} "
+            f"status={snap['status']} optimizer={snap['optimizer_status']} segments={snap['segments_count']}"
+        )
+        time.sleep(max(0.2, float(args.poll_sec)))
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_modify_vectors_smoketest.py

@@ -0,0 +1,29 @@
+import argparse
+from pprint import pprint
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", type=str, required=True)
+    parser.add_argument("--prefer-grpc", action="store_true", default=False)
+    args = parser.parse_args()
+
+    from visual_rag import QdrantAdmin
+
+    admin = QdrantAdmin(prefer_grpc=bool(args.prefer_grpc), timeout=60)
+    before = admin.get_collection_info(collection_name=str(args.collection))
+    print("BEFORE points_count=", before.get("points_count"))
+    existing = sorted(
+        (((before.get("config") or {}).get("params") or {}).get("vectors") or {}).keys()
+    )
+    print("BEFORE vectors=", existing)
+
+    after = admin.ensure_collection_all_on_disk(collection_name=str(args.collection))
+
+    print("AFTER points_count=", after.get("points_count"))
+    print("AFTER params.vectors:")
+    pprint(((after.get("config") or {}).get("params") or {}).get("vectors"))
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_rebuild_collection_no_index.py

@@ -0,0 +1,297 @@
+"""
+Rebuild a Qdrant collection so `indexed_vectors_count` becomes 0 (no ANN/HNSW built).
+
+Important:
+- Qdrant does NOT support "unbuilding" an existing HNSW index in-place.
+- The only reliable way to get indexed_vectors_count back to 0 is to:
+  1) copy points to a temporary collection,
+  2) delete + recreate the original collection with a very large indexing_threshold,
+  3) copy points back,
+  4) delete the temporary collection.
+
+This script keeps the *final* collection name unchanged.
+
+Usage:
+  python scripts/qdrant_rebuild_collection_no_index.py \
+    --collection "my_collection" \
+    --embedding-dim 128 \
+    --vector-dtype float32 \
+    --indexing-threshold 1000000000
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+try:
+    from dotenv import load_dotenv
+
+    DOTENV_AVAILABLE = True
+except Exception:
+    DOTENV_AVAILABLE = False
+
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qm
+
+
+def _get_env(name: str) -> Optional[str]:
+    v = os.getenv(name)
+    if v is None:
+        return None
+    v = str(v).strip()
+    return v or None
+
+
+def _require(value: Optional[str], *, name: str) -> str:
+    if not value:
+        raise SystemExit(f"Missing {name}. Provide flag or set env var.")
+    return value
+
+
+def _chunks(seq: Sequence[Any], n: int) -> List[Sequence[Any]]:
+    return [seq[i : i + n] for i in range(0, len(seq), n)]
+
+
+def _vectors_config(*, embedding_dim: int, vector_dtype: str) -> Dict[str, qm.VectorParams]:
+    datatype = qm.Datatype.FLOAT16 if vector_dtype == "float16" else qm.Datatype.FLOAT32
+    multivector_config = qm.MultiVectorConfig(comparator=qm.MultiVectorComparator.MAX_SIM)
+    return {
+        "initial": qm.VectorParams(
+            size=int(embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=True,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "mean_pooling": qm.VectorParams(
+            size=int(embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "experimental_pooling": qm.VectorParams(
+            size=int(embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            multivector_config=multivector_config,
+            datatype=datatype,
+        ),
+        "global_pooling": qm.VectorParams(
+            size=int(embedding_dim),
+            distance=qm.Distance.COSINE,
+            on_disk=False,
+            datatype=datatype,
+        ),
+    }
+
+
+def _scroll_points(
+    client: QdrantClient,
+    *,
+    collection: str,
+    limit: int,
+    offset: Any,
+) -> Tuple[List[Any], Any]:
+    # qdrant-client returns (points, next_offset)
+    return client.scroll(
+        collection_name=collection,
+        limit=int(limit),
+        with_payload=True,
+        with_vectors=True,
+        offset=offset,
+    )
+
+
+def _clone(
+    client: QdrantClient,
+    *,
+    source: str,
+    dest: str,
+    embedding_dim: int,
+    vector_dtype: str,
+    indexing_threshold: int,
+    recreate_dest: bool,
+    scroll_limit: int,
+    upsert_batch_size: int,
+) -> int:
+    if recreate_dest:
+        try:
+            client.delete_collection(dest)
+        except Exception:
+            pass
+
+    # Create destination collection
+    client.create_collection(
+        collection_name=dest,
+        vectors_config=_vectors_config(embedding_dim=embedding_dim, vector_dtype=vector_dtype),
+        optimizers_config=qm.OptimizersConfigDiff(indexing_threshold=int(indexing_threshold)),
+    )
+    # Keep filename payload index (cheap; useful for skip_existing)
+    try:
+        client.create_payload_index(
+            collection_name=dest,
+            field_name="filename",
+            field_schema=qm.PayloadSchemaType.KEYWORD,
+        )
+    except Exception:
+        pass
+
+    total = 0
+    next_offset = None
+
+    while True:
+        points, next_offset = _scroll_points(
+            client,
+            collection=source,
+            limit=scroll_limit,
+            offset=next_offset,
+        )
+        if not points:
+            break
+
+        for batch in _chunks(points, int(upsert_batch_size)):
+            upsert_points: List[qm.PointStruct] = []
+            for p in batch:
+                upsert_points.append(
+                    qm.PointStruct(
+                        id=p.id,
+                        vector=p.vector,
+                        payload=p.payload or {},
+                    )
+                )
+            client.upsert(collection_name=dest, points=upsert_points, wait=True)
+            total += len(upsert_points)
+            if total % 500 == 0:
+                print(f"… copied {total} points to {dest}")
+
+        if next_offset is None:
+            break
+
+    return total
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--collection", required=True, help="Collection to rebuild (final name stays same)"
+    )
+    parser.add_argument("--qdrant-url", default=None, help="Override QDRANT_URL")
+    parser.add_argument("--qdrant-api-key", default=None, help="Override QDRANT_API_KEY")
+    parser.add_argument("--prefer-grpc", action="store_true", help="Use gRPC transport")
+    parser.add_argument("--timeout", type=float, default=300.0, help="Client timeout seconds")
+    parser.add_argument(
+        "--embedding-dim", type=int, default=128, help="Embedding dim (typically 128)"
+    )
+    parser.add_argument("--vector-dtype", choices=["float16", "float32"], default="float32")
+    parser.add_argument(
+        "--indexing-threshold",
+        type=int,
+        default=1_000_000_000,
+        help="Very large value keeps indexed_vectors_count at 0",
+    )
+    parser.add_argument("--scroll-limit", type=int, default=256)
+    parser.add_argument("--upsert-batch-size", type=int, default=64)
+    parser.add_argument(
+        "--keep-temp", action="store_true", help="Do not delete temp collection at the end"
+    )
+    args = parser.parse_args()
+
+    if DOTENV_AVAILABLE:
+        load_dotenv()
+
+    url = args.qdrant_url or _get_env("QDRANT_URL")
+    api_key = args.qdrant_api_key or _get_env("QDRANT_API_KEY")
+    url = _require(url, name="QDRANT_URL/--qdrant-url")
+    api_key = _require(api_key, name="QDRANT_API_KEY/--qdrant-api-key")
+
+    client = QdrantClient(
+        url=url,
+        api_key=api_key,
+        prefer_grpc=bool(args.prefer_grpc),
+        timeout=float(args.timeout),
+        check_compatibility=False,
+    )
+
+    info = client.get_collection(args.collection)
+    print(f"✅ Found collection: {args.collection}")
+    print(f"   points_count≈{info.points_count} indexed_vectors_count={info.indexed_vectors_count}")
+
+    stamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
+    temp = f"{args.collection}__tmp_rebuild_noindex__{stamp}"
+    print(f"🧪 Temp collection: {temp}")
+
+    print("➡️ Step 1/4: Copying points to temp…")
+    copied1 = _clone(
+        client,
+        source=args.collection,
+        dest=temp,
+        embedding_dim=int(args.embedding_dim),
+        vector_dtype=str(args.vector_dtype),
+        indexing_threshold=int(args.indexing_threshold),
+        recreate_dest=True,
+        scroll_limit=int(args.scroll_limit),
+        upsert_batch_size=int(args.upsert_batch_size),
+    )
+    temp_info = client.get_collection(temp)
+    print(
+        f"✅ Temp ready: points_count≈{temp_info.points_count} indexed_vectors_count={temp_info.indexed_vectors_count}"
+    )
+
+    print("➡️ Step 2/4: Deleting original collection…")
+    client.delete_collection(args.collection)
+    time.sleep(1.0)
+
+    print("➡️ Step 3/4: Recreating original with indexing disabled…")
+    client.create_collection(
+        collection_name=args.collection,
+        vectors_config=_vectors_config(
+            embedding_dim=int(args.embedding_dim), vector_dtype=str(args.vector_dtype)
+        ),
+        optimizers_config=qm.OptimizersConfigDiff(indexing_threshold=int(args.indexing_threshold)),
+    )
+    try:
+        client.create_payload_index(
+            collection_name=args.collection,
+            field_name="filename",
+            field_schema=qm.PayloadSchemaType.KEYWORD,
+        )
+    except Exception:
+        pass
+
+    print("➡️ Step 4/4: Copying points back to original…")
+    copied2 = _clone(
+        client,
+        source=temp,
+        dest=args.collection,
+        embedding_dim=int(args.embedding_dim),
+        vector_dtype=str(args.vector_dtype),
+        indexing_threshold=int(args.indexing_threshold),
+        recreate_dest=False,
+        scroll_limit=int(args.scroll_limit),
+        upsert_batch_size=int(args.upsert_batch_size),
+    )
+
+    final_info = client.get_collection(args.collection)
+    exact = client.count(collection_name=args.collection, exact=True)
+    print("✅ Rebuild complete")
+    print(f"   copied_to_temp={copied1} copied_back={copied2}")
+    print(
+        f"   final.points_count≈{final_info.points_count} final.count(exact)={exact.count} "
+        f"final.indexed_vectors_count={final_info.indexed_vectors_count}"
+    )
+
+    if not args.keep_temp:
+        print("🧹 Deleting temp collection…")
+        client.delete_collection(temp)
+        print("✅ Temp deleted")
+    else:
+        print("ℹ️ Temp kept:", temp)
+
+
+if __name__ == "__main__":
+    main()

scripts/qdrant_recompute_colqwen_pooling_from_initial.py

@@ -0,0 +1,312 @@
+"""
+Recompute ColQwen2.5 pooled vectors from already-indexed `initial` vectors.
+
+Why:
+- Your collection already contains high-quality `initial` multi-vectors (single_full works well),
+  but two-stage prefetch using `experimental_pooling` is poor.
+- We can fix that WITHOUT re-indexing images by recomputing:
+  - mean_pooling  (32×dim)  from `initial` (H×W×dim)
+  - experimental_pooling (36×dim) from mean_pooling with window=5
+  - global_pooling (dim)
+
+How we infer (H, W):
+- For each point we know `num_tokens=len(initial)` and the stored resized image aspect ratio.
+- We factor `num_tokens` and pick the factor pair (h, w) whose w/h best matches width/height.
+
+Usage:
+  python scripts/qdrant_recompute_colqwen_pooling_from_initial.py \
+    --collection "vidore_beir_v2_3ds__colqwen25_v0_2__nocrop__union__fp32__grpc" \
+    --dataset "vidore/esg_reports_v2" \
+    --limit 0
+"""
+
+from __future__ import annotations
+
+import argparse
+import math
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import numpy as np
+
+try:
+    from dotenv import load_dotenv
+
+    DOTENV_AVAILABLE = True
+except Exception:
+    DOTENV_AVAILABLE = False
+
+from qdrant_client import QdrantClient
+from qdrant_client.http import models as qm
+
+from visual_rag.embedding.pooling import (
+    adaptive_row_mean_pooling_from_grid,
+    colpali_experimental_pooling_from_rows,
+)
+
+
+def _maybe_load_dotenv() -> None:
+    if not DOTENV_AVAILABLE:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _stable_uuid(text: str) -> str:
+    import hashlib
+
+    hex_str = hashlib.sha256(text.encode("utf-8")).hexdigest()[:32]
+    return f"{hex_str[:8]}-{hex_str[8:12]}-{hex_str[12:16]}-{hex_str[16:20]}-{hex_str[20:32]}"
+
+
+def _infer_grid(num_tokens: int, *, width: Optional[int], height: Optional[int]) -> Tuple[int, int]:
+    """
+    Infer (grid_h, grid_w) such that grid_h*grid_w=num_tokens.
+
+    Picks factor pair closest to the observed aspect ratio (width/height).
+    """
+    n = int(num_tokens)
+    if n <= 0:
+        raise ValueError("num_tokens must be > 0")
+
+    # Fallback aspect if missing
+    if width and height and int(width) > 0 and int(height) > 0:
+        aspect = float(width) / float(height)
+    else:
+        aspect = 1.0
+
+    best = None
+    best_score = float("inf")
+
+    # Enumerate factors up to sqrt(n)
+    r = int(math.isqrt(n))
+    for h in range(1, r + 1):
+        if n % h != 0:
+            continue
+        w = n // h
+
+        # Consider both orientations
+        for hh, ww in ((h, w), (w, h)):
+            if hh <= 0 or ww <= 0:
+                continue
+            cand = float(ww) / float(hh)
+            # log-space ratio distance is symmetric and scale-invariant
+            score = abs(math.log(max(cand, 1e-9) / max(aspect, 1e-9)))
+            if score < best_score:
+                best_score = score
+                best = (int(hh), int(ww))
+
+    if best is None:
+        # Should never happen
+        g = int(round(math.sqrt(n)))
+        return g, max(1, n // max(1, g))
+    return best
+
+
+def _chunks(xs: List[Any], n: int) -> Iterable[List[Any]]:
+    for i in range(0, len(xs), n):
+        yield xs[i : i + n]
+
+
+def _has_none_nested(v: Any) -> bool:
+    try:
+        if not isinstance(v, list):
+            return True
+        if not v:
+            return True
+        if not isinstance(v[0], list):
+            return True
+        for row in v:
+            if not isinstance(row, list):
+                return True
+            for x in row:
+                if x is None:
+                    return True
+        return False
+    except Exception:
+        return True
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--collection", required=True)
+    ap.add_argument("--dataset", required=True, help="payload['dataset'] value to filter on")
+    ap.add_argument("--url", default="")
+    ap.add_argument("--api-key", default="")
+    ap.add_argument("--timeout", type=float, default=120.0)
+    ap.add_argument("--scroll-limit", type=int, default=128)
+    ap.add_argument("--update-batch", type=int, default=64)
+    ap.add_argument(
+        "--retrieve-batch", type=int, default=16, help="Batch size for retrieve() calls"
+    )
+    ap.add_argument("--limit", type=int, default=0, help="0 means no limit")
+    ap.add_argument("--sleep-sec", type=float, default=0.0)
+    args = ap.parse_args()
+
+    _maybe_load_dotenv()
+
+    url = args.url or os.getenv("QDRANT_URL") or ""
+    if not url:
+        raise SystemExit("QDRANT_URL not set (or pass --url)")
+    api_key = args.api_key or os.getenv("QDRANT_API_KEY") or None
+
+    client = QdrantClient(
+        url=url,
+        api_key=api_key,
+        prefer_grpc=False,  # avoid DNS issues for 6334 in some envs
+        timeout=float(args.timeout),
+        check_compatibility=False,
+    )
+
+    flt = qm.Filter(
+        must=[qm.FieldCondition(key="dataset", match=qm.MatchValue(value=str(args.dataset)))]
+    )
+
+    updated = 0
+    scanned = 0
+    next_offset = None
+
+    while True:
+        points, next_offset = client.scroll(
+            collection_name=str(args.collection),
+            scroll_filter=flt,
+            limit=int(args.scroll_limit),
+            offset=next_offset,
+            with_payload=True,
+            with_vectors=False,  # retrieve vectors per-point to avoid whole-batch parse failures
+        )
+        if not points:
+            break
+
+        pv_batch: List[qm.PointVectors] = []
+        ids: List[Any] = [p.id for p in points]
+        payload_by_id: Dict[Any, Dict[str, Any]] = {p.id: (p.payload or {}) for p in points}
+
+        # Retrieve initial vectors in batches for speed; fallback to per-id retrieve on failure.
+        records_by_id: Dict[Any, Any] = {}
+        for id_chunk in _chunks(ids, int(args.retrieve_batch)):
+            if not id_chunk:
+                continue
+            try:
+                recs = client.retrieve(
+                    collection_name=str(args.collection),
+                    ids=id_chunk,
+                    with_payload=False,
+                    with_vectors=["initial"],
+                    timeout=int(args.timeout),
+                )
+                for r in recs:
+                    records_by_id[r.id] = r
+            except Exception:
+                # fallback: per-id
+                for pid in id_chunk:
+                    try:
+                        recs = client.retrieve(
+                            collection_name=str(args.collection),
+                            ids=[pid],
+                            with_payload=False,
+                            with_vectors=["initial"],
+                            timeout=int(args.timeout),
+                        )
+                        if recs:
+                            records_by_id[recs[0].id] = recs[0]
+                    except Exception:
+                        continue
+
+        for pid in ids:
+            scanned += 1
+            if args.limit and scanned > int(args.limit):
+                break
+
+            # Retrieve vectors for this point. Some points in this collection may contain placeholder
+            # vectors with nulls from recovery attempts; retrieving per-id lets us skip them safely.
+            rec = records_by_id.get(pid)
+            if rec is None:
+                continue
+            vec = (rec.vector or {}).get("initial")
+            if _has_none_nested(vec):
+                continue
+
+            emb = np.asarray(vec, dtype=np.float32)  # [num_tokens, dim]
+            num_tokens = int(emb.shape[0])
+
+            payload = payload_by_id.get(pid) or {}
+            w = (
+                payload.get("resized_width")
+                or payload.get("cropped_width")
+                or payload.get("original_width")
+            )
+            h = (
+                payload.get("resized_height")
+                or payload.get("cropped_height")
+                or payload.get("original_height")
+            )
+            try:
+                w_i = int(w) if w is not None else None
+                h_i = int(h) if h is not None else None
+            except Exception:
+                w_i, h_i = None, None
+
+            grid_h, grid_w = _infer_grid(num_tokens, width=w_i, height=h_i)
+            if grid_h * grid_w != num_tokens:
+                # safety: if factor inference failed, skip
+                continue
+
+            mean_pool = adaptive_row_mean_pooling_from_grid(
+                emb,
+                grid_h=int(grid_h),
+                grid_w=int(grid_w),
+                target_rows=32,  # IMPORTANT: fixed row count for good prefetch recall
+                output_dtype=np.float32,
+            )
+            exp_pool = colpali_experimental_pooling_from_rows(
+                mean_pool,
+                window_size=5,
+                output_dtype=np.float32,
+            )
+            glob = mean_pool.mean(axis=0).astype(np.float32)
+
+            pv_batch.append(
+                qm.PointVectors(
+                    id=pid,
+                    vector={
+                        "mean_pooling": mean_pool.tolist(),
+                        "experimental_pooling": exp_pool.tolist(),
+                        "global_pooling": glob.tolist(),
+                    },
+                )
+            )
+
+            if len(pv_batch) >= int(args.update_batch):
+                client.update_vectors(
+                    collection_name=str(args.collection),
+                    points=pv_batch,
+                    wait=True,
+                )
+                updated += len(pv_batch)
+                print(f"✅ updated vectors: {updated} (scanned={scanned})", flush=True)
+                pv_batch = []
+                if float(args.sleep_sec) > 0:
+                    time.sleep(float(args.sleep_sec))
+
+        if pv_batch:
+            client.update_vectors(
+                collection_name=str(args.collection),
+                points=pv_batch,
+                wait=True,
+            )
+            updated += len(pv_batch)
+            print(f"✅ updated vectors: {updated} (scanned={scanned})", flush=True)
+
+        if args.limit and scanned >= int(args.limit):
+            break
+        if next_offset is None:
+            break
+
+    print(f"Done. scanned={scanned}, updated={updated}")
+
+
+if __name__ == "__main__":
+    main()

scripts/query_token_stats.py

@@ -0,0 +1,123 @@
+import argparse
+import json
+from pathlib import Path
+from typing import Any, Dict, List
+
+
+def _stats(xs: List[int]) -> Dict[str, Any]:
+    import numpy as np
+
+    if not xs:
+        return {
+            "count": 0,
+            "mean": None,
+            "median": None,
+            "p95": None,
+            "min": None,
+            "max": None,
+        }
+    arr = np.array(xs, dtype=np.int64)
+    return {
+        "count": int(arr.size),
+        "mean": float(arr.mean()),
+        "median": float(np.median(arr)),
+        "p95": float(np.percentile(arr, 95)),
+        "min": int(arr.min()),
+        "max": int(arr.max()),
+    }
+
+
+def _count_tokens(emb) -> int:
+    try:
+        import torch
+
+        if isinstance(emb, torch.Tensor):
+            return int(emb.shape[0])
+    except Exception:
+        pass
+    try:
+        return int(emb.shape[0])
+    except Exception:
+        return int(len(emb))
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", type=str, default=None)
+    parser.add_argument("--datasets", type=str, nargs="+", default=None)
+    parser.add_argument("--model", type=str, default="vidore/colSmol-500M")
+    parser.add_argument(
+        "--torch-dtype",
+        type=str,
+        default="auto",
+        choices=["auto", "float32", "float16", "bfloat16"],
+    )
+    parser.add_argument(
+        "--processor-speed", type=str, default="fast", choices=["fast", "slow", "auto"]
+    )
+    parser.add_argument("--batch-size", type=int, default=16)
+    parser.add_argument("--no-filter-special-tokens", action="store_true", default=False)
+    parser.add_argument("--max-queries", type=int, default=0)
+    parser.add_argument("--output", type=str, default="")
+    args = parser.parse_args()
+
+    datasets: List[str] = []
+    if args.datasets:
+        datasets = list(args.datasets)
+    elif args.dataset:
+        datasets = [args.dataset]
+    else:
+        raise SystemExit("Provide --dataset or --datasets")
+
+    from benchmarks.vidore_beir_qdrant.run_qdrant_beir import _maybe_load_dotenv, _parse_torch_dtype
+    from benchmarks.vidore_tatdqa_test.dataset_loader import load_vidore_beir_dataset
+    from visual_rag.embedding.visual_embedder import VisualEmbedder
+
+    _maybe_load_dotenv()
+
+    embedder = VisualEmbedder(
+        model_name=str(args.model),
+        torch_dtype=_parse_torch_dtype(str(args.torch_dtype)),
+        batch_size=int(args.batch_size),
+        processor_speed=str(args.processor_speed),
+    )
+    filter_special = not bool(args.no_filter_special_tokens)
+
+    out: Dict[str, Any] = {
+        "model": str(args.model),
+        "torch_dtype": str(args.torch_dtype),
+        "processor_speed": str(args.processor_speed),
+        "filter_special_tokens": bool(filter_special),
+        "max_queries": int(args.max_queries),
+        "datasets": {},
+    }
+
+    for ds in datasets:
+        _, queries, _ = load_vidore_beir_dataset(ds)
+        qs = [q.text for q in queries]
+        if int(args.max_queries) and int(args.max_queries) > 0:
+            qs = qs[: int(args.max_queries)]
+        embs = embedder.embed_queries(
+            qs,
+            batch_size=int(args.batch_size),
+            filter_special_tokens=bool(filter_special),
+            show_progress=True,
+        )
+        token_counts = [_count_tokens(e) for e in embs]
+        out["datasets"][str(ds)] = {
+            "num_queries": int(len(qs)),
+            "token_count": _stats(token_counts),
+        }
+
+    text = json.dumps(out, indent=2)
+    if args.output:
+        p = Path(str(args.output))
+        p.parent.mkdir(parents=True, exist_ok=True)
+        p.write_text(text, encoding="utf-8")
+        print(str(p))
+    else:
+        print(text)
+
+
+if __name__ == "__main__":
+    main()

scripts/update_qdrant_indexing_threshold.py

@@ -0,0 +1,171 @@
+import argparse
+import json
+import os
+import time
+from pathlib import Path
+
+
+def _maybe_load_dotenv() -> None:
+    try:
+        from dotenv import load_dotenv
+    except ImportError:
+        return
+    if Path(".env").exists():
+        load_dotenv(".env")
+
+
+def _as_jsonable(obj):
+    if obj is None:
+        return None
+    if isinstance(obj, (str, int, float, bool)):
+        return obj
+    if isinstance(obj, dict):
+        return {str(k): _as_jsonable(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return [_as_jsonable(v) for v in obj]
+    if hasattr(obj, "model_dump"):
+        try:
+            return obj.model_dump()
+        except Exception:
+            pass
+    if hasattr(obj, "__dict__"):
+        try:
+            return {k: _as_jsonable(v) for k, v in obj.__dict__.items()}
+        except Exception:
+            pass
+    return str(obj)
+
+
+def _indexed_total(indexed_vectors_count) -> int:
+    if indexed_vectors_count is None:
+        return 0
+    if isinstance(indexed_vectors_count, dict):
+        try:
+            return int(sum(int(v) for v in indexed_vectors_count.values()))
+        except Exception:
+            return 0
+    try:
+        return int(indexed_vectors_count)
+    except Exception:
+        return 0
+
+
+def _snapshot_info(info) -> dict:
+    status = getattr(info, "status", None)
+    if hasattr(status, "value"):
+        status = status.value
+    optimizer_status = getattr(info, "optimizer_status", None)
+    if hasattr(optimizer_status, "value"):
+        optimizer_status = optimizer_status.value
+    return {
+        "status": _as_jsonable(status),
+        "optimizer_status": _as_jsonable(optimizer_status),
+        "points_count": _as_jsonable(getattr(info, "points_count", None)),
+        "indexed_vectors_count": _as_jsonable(getattr(info, "indexed_vectors_count", None)),
+        "segments_count": _as_jsonable(getattr(info, "segments_count", None)),
+    }
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--collection", type=str, required=True)
+    parser.add_argument("--indexing-threshold", type=int, default=0)
+    parser.add_argument("--prefer-grpc", action="store_true")
+    parser.add_argument("--url", type=str, default="")
+    parser.add_argument("--api-key", type=str, default="")
+    parser.add_argument("--wait", action="store_true")
+    parser.add_argument("--timeout-sec", type=int, default=300)
+    parser.add_argument("--poll-sec", type=int, default=2)
+    parser.add_argument("--dump-json", type=str, default="")
+    args = parser.parse_args()
+
+    _maybe_load_dotenv()
+
+    qdrant_url = args.url or os.getenv("QDRANT_URL")
+    if not qdrant_url:
+        raise ValueError("QDRANT_URL not set")
+    qdrant_api_key = args.api_key or os.getenv("QDRANT_API_KEY")
+
+    from qdrant_client import QdrantClient
+    from qdrant_client.http import models
+
+    client = QdrantClient(
+        url=qdrant_url,
+        api_key=qdrant_api_key,
+        prefer_grpc=args.prefer_grpc,
+        check_compatibility=False,
+        timeout=60,
+    )
+
+    client.update_collection(
+        collection_name=args.collection,
+        optimizers_config=models.OptimizersConfigDiff(
+            indexing_threshold=int(args.indexing_threshold)
+        ),
+    )
+
+    info = client.get_collection(args.collection)
+    snap = _snapshot_info(info)
+    print(
+        f"Updated optimizers.indexing_threshold={args.indexing_threshold} for collection='{args.collection}'. "
+        f"points={snap['points_count']}, indexed_vectors={snap['indexed_vectors_count']}, "
+        f"status={snap['status']}, optimizer_status={snap['optimizer_status']}, segments={snap['segments_count']}"
+    )
+    if args.dump_json:
+        out_path = Path(args.dump_json)
+        out_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(out_path, "w") as f:
+            json.dump(
+                {"event": "after_update", "collection": args.collection, "snapshot": snap},
+                f,
+                indent=2,
+            )
+
+    if not args.wait:
+        return
+
+    start = time.time()
+    while True:
+        info = client.get_collection(args.collection)
+        snap = _snapshot_info(info)
+        indexed_total = _indexed_total(snap["indexed_vectors_count"])
+        total = int(snap["points_count"] or 0)
+        if indexed_total >= total and total > 0:
+            print(
+                f"Indexing complete: indexed_vectors={snap['indexed_vectors_count']}, points={snap['points_count']}"
+            )
+            if args.dump_json:
+                out_path = Path(args.dump_json)
+                out_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(out_path, "w") as f:
+                    json.dump(
+                        {"event": "complete", "collection": args.collection, "snapshot": snap},
+                        f,
+                        indent=2,
+                    )
+            return
+        if time.time() - start > args.timeout_sec:
+            print(
+                f"Timeout waiting for indexing: indexed_vectors={snap['indexed_vectors_count']}, "
+                f"points={snap['points_count']}, status={snap['status']}, optimizer_status={snap['optimizer_status']}"
+            )
+            if args.dump_json:
+                out_path = Path(args.dump_json)
+                out_path.parent.mkdir(parents=True, exist_ok=True)
+                with open(out_path, "w") as f:
+                    json.dump(
+                        {"event": "timeout", "collection": args.collection, "snapshot": snap},
+                        f,
+                        indent=2,
+                    )
+            return
+        print(
+            f"Indexing in progress: indexed_vectors={snap['indexed_vectors_count']}, "
+            f"points={snap['points_count']}, status={snap['status']}, optimizer_status={snap['optimizer_status']}, "
+            f"segments={snap['segments_count']}"
+        )
+        time.sleep(max(0.1, float(args.poll_sec)))
+
+
+if __name__ == "__main__":
+    main()

tests/__init__.py

@@ -1,8 +1 @@
 # Tests for visual-rag-toolkit
-
-
-
-
-
-
-

tests/test_config.py

@@ -1,18 +1,16 @@
 """Tests for configuration utilities."""
 
 import os
-import pytest
 import tempfile
-from pathlib import Path
 
 
 class TestConfigLoading:
     """Test config file loading."""
-    
+
     def test_load_yaml_config(self):
         """Load config from YAML file."""
         from visual_rag.config import load_config
-        
+
         # Create temp config file
         config_content = """
 model:
@@ -21,34 +19,34 @@ model:
 qdrant:
   url: "http://localhost:6333"
 """
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
             f.write(config_content)
             config_path = f.name
-        
+
         try:
             config = load_config(config_path, force_reload=True, apply_env_overrides=False)
-            
+
             assert config["model"]["name"] == "test-model"
             assert config["model"]["batch_size"] == 8
             assert config["qdrant"]["url"] == "http://localhost:6333"
         finally:
             os.unlink(config_path)
-    
+
     def test_env_override(self):
         """Environment variables override config values."""
-        from visual_rag.config import load_config, get
-        
+        from visual_rag.config import load_config
+
         # Set env var
         os.environ["VISUAL_RAG_MODEL_NAME"] = "env-override-model"
-        
+
         config_content = """
 model:
   name: "yaml-model"
 """
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
             f.write(config_content)
             config_path = f.name
-        
+
         try:
             config = load_config(config_path, force_reload=True, apply_env_overrides=False)
             # The env var should be checked in get() if implemented
@@ -57,20 +55,19 @@ model:
         finally:
             os.unlink(config_path)
             del os.environ["VISUAL_RAG_MODEL_NAME"]
-    
+
     def test_missing_config_uses_defaults(self):
         """Missing config file returns empty dict or defaults."""
         from visual_rag.config import load_config
-        
+
         config = load_config("/nonexistent/path/config.yaml")
-        
+
         # Should not raise, returns empty or default config
         assert isinstance(config, dict)
-    
+
     def test_get_nested_value(self):
         """Get nested config values with dot notation."""
-        from visual_rag.config import get
-        
+
         # This tests the get() function if available
         # Will need the config to be loaded first
         pass  # Placeholder - depends on implementation
@@ -78,39 +75,32 @@ model:
 
 class TestConfigSection:
     """Test getting config sections."""
-    
+
     def test_get_section(self):
         """Get a config section."""
         from visual_rag.config import get_section, load_config
-        
+
         config_content = """
 qdrant:
   url: "http://localhost"
   collection: "test"
 """
-        with tempfile.NamedTemporaryFile(mode='w', suffix='.yaml', delete=False) as f:
+        with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f:
             f.write(config_content)
             config_path = f.name
-        
+
         try:
             load_config(config_path, force_reload=True, apply_env_overrides=False)
             section = get_section("qdrant", apply_env_overrides=False)
-            
+
             assert section["url"] == "http://localhost"
             assert section["collection"] == "test"
         finally:
             os.unlink(config_path)
-    
+
     def test_missing_section(self):
         """Missing section returns empty dict."""
         from visual_rag.config import get_section
-        
+
         section = get_section("nonexistent")
         assert section == {}
-
-
-
-
-
-
-

tests/test_pdf_processor.py

@@ -1,88 +1,87 @@
 """Tests for PDF processor."""
 
 import pytest
-import numpy as np
 from PIL import Image
 
 
 class TestResizeForColPali:
     """Test image resizing for ColPali processing."""
-    
+
     def test_resize_standard_image(self):
         """Standard image resizes to tile boundaries."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        
+
         # Create test image (A4-like ratio)
-        img = Image.new('RGB', (1000, 1414), color='white')
-        
+        img = Image.new("RGB", (1000, 1414), color="white")
+
         resized, tile_rows, tile_cols = processor.resize_for_colpali(img)
-        
+
         # Should resize to multiples of 512
         assert resized.width % 512 == 0 or resized.width <= 2048
         assert resized.height % 512 == 0 or resized.height <= 2048
         assert tile_rows >= 1
         assert tile_cols >= 1
-    
+
     def test_resize_small_image(self):
         """Small image handles gracefully."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        img = Image.new('RGB', (100, 100), color='white')
-        
+        img = Image.new("RGB", (100, 100), color="white")
+
         resized, tile_rows, tile_cols = processor.resize_for_colpali(img)
-        
+
         assert resized is not None
         assert tile_rows >= 1
         assert tile_cols >= 1
-    
+
     def test_resize_wide_image(self):
         """Wide image (panorama-like) resizes correctly."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        img = Image.new('RGB', (3000, 500), color='white')
-        
+        img = Image.new("RGB", (3000, 500), color="white")
+
         resized, tile_rows, tile_cols = processor.resize_for_colpali(img)
-        
+
         # Wide image should have more cols than rows
         assert tile_cols >= tile_rows
-    
+
     def test_resize_preserves_rgb(self):
         """Resized image should be RGB."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        img = Image.new('RGBA', (1000, 1000), color='white')
-        
+        img = Image.new("RGBA", (1000, 1000), color="white")
+
         resized, _, _ = processor.resize_for_colpali(img)
-        
-        assert resized.mode == 'RGB'
+
+        assert resized.mode == "RGB"
 
 
 class TestMetadataExtraction:
     """Test metadata extraction from filenames."""
-    
+
     def test_extract_year_from_filename(self):
         """Extract year from filename."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        
+
         metadata = processor.extract_metadata_from_filename("Annual_Report_2023.pdf")
-        
+
         # Should extract year if implemented
         # This depends on your implementation
         assert isinstance(metadata, dict)
-    
+
     def test_sanitize_filename(self):
         """Sanitize filename for safe storage."""
         from visual_rag.indexing.pdf_processor import PDFProcessor
-        
+
         processor = PDFProcessor()
-        
+
         # Test with special characters
         filename = "Report (Final) - v2.0.pdf"
         # Should handle gracefully
@@ -92,41 +91,35 @@ class TestMetadataExtraction:
 
 class TestChunkIdGeneration:
     """Test deterministic chunk ID generation."""
-    
+
     def test_chunk_id_deterministic(self):
         """Same input produces same chunk ID."""
         from visual_rag.indexing.pipeline import ProcessingPipeline
-        
+
         id1 = ProcessingPipeline.generate_chunk_id("test.pdf", 1)
         id2 = ProcessingPipeline.generate_chunk_id("test.pdf", 1)
-        
+
         assert id1 == id2
-    
+
     def test_chunk_id_unique(self):
         """Different pages produce different IDs."""
         from visual_rag.indexing.pipeline import ProcessingPipeline
-        
+
         id1 = ProcessingPipeline.generate_chunk_id("test.pdf", 1)
         id2 = ProcessingPipeline.generate_chunk_id("test.pdf", 2)
-        
+
         assert id1 != id2
-    
+
     def test_chunk_id_format(self):
         """Chunk ID should be valid UUID format."""
-        from visual_rag.indexing.pipeline import ProcessingPipeline
         import uuid
-        
+
+        from visual_rag.indexing.pipeline import ProcessingPipeline
+
         chunk_id = ProcessingPipeline.generate_chunk_id("test.pdf", 1)
-        
+
         # Should be valid UUID
         try:
             uuid.UUID(chunk_id)
         except ValueError:
             pytest.fail(f"Invalid UUID format: {chunk_id}")
-
-
-
-
-
-
-

tests/test_pooling.py

@@ -1,55 +1,54 @@
 """Tests for pooling functions."""
 
-import pytest
 import numpy as np
 
 
 class TestTileLevelPooling:
     """Test tile-level mean pooling."""
-    
+
     def test_basic_pooling(self):
         """Pooling reduces [num_tokens, dim] → [num_tiles, dim]."""
         from visual_rag.embedding.pooling import tile_level_mean_pooling
-        
+
         # 13 tiles × 64 patches = 832 visual tokens
         num_tiles = 13
         patches_per_tile = 64
         num_tokens = num_tiles * patches_per_tile
         dim = 128
-        
+
         embedding = np.random.randn(num_tokens, dim).astype(np.float32)
         pooled = tile_level_mean_pooling(embedding, num_tiles, patches_per_tile)
-        
+
         assert pooled.shape == (num_tiles, dim)
         assert pooled.dtype == np.float32
-    
+
     def test_pooling_preserves_info(self):
         """Pooled vectors should be mean of patches."""
         from visual_rag.embedding.pooling import tile_level_mean_pooling
-        
+
         num_tiles = 5
         patches_per_tile = 64
         dim = 128
-        
+
         embedding = np.random.randn(num_tiles * patches_per_tile, dim).astype(np.float32)
         pooled = tile_level_mean_pooling(embedding, num_tiles, patches_per_tile)
-        
+
         # Check first tile
         expected_tile0 = embedding[:patches_per_tile].mean(axis=0)
         np.testing.assert_array_almost_equal(pooled[0], expected_tile0, decimal=5)
-    
+
     def test_pooling_with_partial_last_tile(self):
         """Handle case where last tile has fewer patches."""
         from visual_rag.embedding.pooling import tile_level_mean_pooling
-        
+
         # 800 tokens, 64 per tile = 12.5 tiles → 13 tiles with partial last
         num_tokens = 800
         num_tiles = 13
         dim = 128
-        
+
         embedding = np.random.randn(num_tokens, dim).astype(np.float32)
         pooled = tile_level_mean_pooling(embedding, num_tiles, patches_per_tile=64)
-        
+
         # Should handle gracefully - at least some tiles
         assert pooled.shape[1] == dim
         assert pooled.shape[0] >= 1
@@ -57,14 +56,14 @@ class TestTileLevelPooling:
 
 class TestGlobalPooling:
     """Test global mean pooling."""
-    
+
     def test_global_mean(self):
         """Global pooling reduces to single vector."""
         from visual_rag.embedding.pooling import global_mean_pooling
-        
+
         embedding = np.random.randn(832, 128).astype(np.float32)
         pooled = global_mean_pooling(embedding)
-        
+
         assert pooled.shape == (128,)
         np.testing.assert_array_almost_equal(pooled, embedding.mean(axis=0))
 
@@ -160,40 +159,40 @@ class TestColPaliExperimentalPooling:
 
 class TestMaxSimScore:
     """Test MaxSim scoring."""
-    
+
     def test_maxsim_identical(self):
         """Identical embeddings should have high score."""
         from visual_rag.embedding.pooling import compute_maxsim_score
-        
+
         embedding = np.random.randn(10, 128).astype(np.float32)
         # Normalize
         embedding = embedding / np.linalg.norm(embedding, axis=1, keepdims=True)
-        
+
         score = compute_maxsim_score(embedding, embedding)
-        
+
         # Should be close to num_tokens (each token matches itself perfectly)
         assert score >= 9.0  # Allow some floating point tolerance
-    
+
     def test_maxsim_orthogonal(self):
         """Orthogonal embeddings should have low score."""
         from visual_rag.embedding.pooling import compute_maxsim_score
-        
+
         # Create orthogonal vectors
         query = np.array([[1, 0, 0, 0], [0, 1, 0, 0]], dtype=np.float32)
         doc = np.array([[0, 0, 1, 0], [0, 0, 0, 1]], dtype=np.float32)
-        
+
         score = compute_maxsim_score(query, doc)
-        
+
         assert score < 0.1  # Near zero for orthogonal
-    
+
     def test_maxsim_shape_independence(self):
         """Score should work with different query/doc lengths."""
         from visual_rag.embedding.pooling import compute_maxsim_score
-        
+
         query = np.random.randn(5, 128).astype(np.float32)
         doc = np.random.randn(100, 128).astype(np.float32)
-        
+
         score = compute_maxsim_score(query, doc)
-        
+
         assert isinstance(score, float)
         assert not np.isnan(score)