diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py index c0d899685f477ab5a5bdaf7d19c40bb3ef63223b..5c08f6f3a8e675075cff170f0498d38e0289e02a 100644 --- a/benchmarks/__init__.py +++ b/benchmarks/__init__.py @@ -8,4 +8,3 @@ work in Docker/Spaces environments. """ __all__ = [] - diff --git a/benchmarks/analyze_results.py b/benchmarks/analyze_results.py index e6fdd592d83fbfee5dfa400b0f2cf80038ccebbd..abc8456445120eb7f4966b1fe4d5aa4fc88336f8 100644 --- a/benchmarks/analyze_results.py +++ b/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() - - - - - - - diff --git a/benchmarks/prepare_submission.py b/benchmarks/prepare_submission.py index c569c3a6323ad803b27f2f97ff37c40867ba4193..ab7c4b8e7e7df469d0dbd0ec04cfa7a7bd8055cd 100644 --- a/benchmarks/prepare_submission.py +++ b/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() - - - - - - - diff --git a/benchmarks/quick_test.py b/benchmarks/quick_test.py index bc633e093b429a0536026522344bc85e8f3afcbf..a9079333666c59c510045d364dc50e975458f754 100644 --- a/benchmarks/quick_test.py +++ b/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) diff --git a/benchmarks/run_vidore.py b/benchmarks/run_vidore.py index c4e42730b2cdeeb8e47ebffe064354f3c0f1f9b0..1f42e842a4f058c5a32bcd99b68b0180f5efe3d9 100644 --- a/benchmarks/run_vidore.py +++ b/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() - - diff --git a/benchmarks/vidore_beir_qdrant/run_qdrant_beir.py b/benchmarks/vidore_beir_qdrant/run_qdrant_beir.py index e733a38604b996b6d0e461bf05806ae9a51919dc..943b13292c20b70b46a20d67a53986cc9d11a955 100644 --- a/benchmarks/vidore_beir_qdrant/run_qdrant_beir.py +++ b/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____.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() \ No newline at end of file + main() diff --git a/benchmarks/vidore_tatdqa_test/__init__.py b/benchmarks/vidore_tatdqa_test/__init__.py index 9b15a1e7e9ab666823f4eeae9598e734148cca5a..a9a2c5b3bb437bff74e283b62c894075e8c15331 100644 --- a/benchmarks/vidore_tatdqa_test/__init__.py +++ b/benchmarks/vidore_tatdqa_test/__init__.py @@ -1,6 +1 @@ __all__ = [] - - - - - diff --git a/benchmarks/vidore_tatdqa_test/dataset_loader.py b/benchmarks/vidore_tatdqa_test/dataset_loader.py index 2a0454eee41f22fc59fc7d982d8e87a148d293d1..cd46ffceae8d15a3d77e62b19808a75406bcb9da 100644 --- a/benchmarks/vidore_tatdqa_test/dataset_loader.py +++ b/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" - - diff --git a/benchmarks/vidore_tatdqa_test/metrics.py b/benchmarks/vidore_tatdqa_test/metrics.py index 0a6472fafc9383701a0dfed252db9f998c1ca3e9..3ee06aa19dfec6c42ccb2d9902b6eb3ded66bc66 100644 --- a/benchmarks/vidore_tatdqa_test/metrics.py +++ b/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) - - - - - diff --git a/benchmarks/vidore_tatdqa_test/run_qdrant.py b/benchmarks/vidore_tatdqa_test/run_qdrant.py index 64fd6832bc8d29a17cfa4a9c0bb8418539eb3888..f27f226a7ab71b9b42179a81607237b9ab6d7630 100644 --- a/benchmarks/vidore_tatdqa_test/run_qdrant.py +++ b/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() - - diff --git a/benchmarks/vidore_tatdqa_test/sweep_eval.py b/benchmarks/vidore_tatdqa_test/sweep_eval.py index 49b933e100ff129f8d0a8dc7269d5b9dfefaf3a6..994b40b70923bbe59b38eab3606b0e7178b9d46b 100644 --- a/benchmarks/vidore_tatdqa_test/sweep_eval.py +++ b/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() - - diff --git a/demo/app.py b/demo/app.py index 4d2d98adf5e60065ae6857fb717a2d991efbe33b..3c8b6ba6b0847656af8ee44476e40e9a6392d0c0 100644 --- a/demo/app.py +++ b/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(): diff --git a/demo/commands.py b/demo/commands.py index 2c0d96d07766a6c67730ac98c0dc2376dcc08729..8bdec130fce1499ec3b1c348b4b8437b0f7096a5 100644 --- a/demo/commands.py +++ b/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) diff --git a/demo/config.py b/demo/config.py index df3fa71a7ba25b22e63dba56162d9fc75ec5903e..a848fed489b7fb4009bac04729d3137e09ceb3d1 100644 --- a/demo/config.py +++ b/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", ] diff --git a/demo/download_models.py b/demo/download_models.py index 8ea655ace9ee56b8f7cd547b1752f713a6a8a7c4..aeb5c3312b5b3f43a96523e503c08b1d54047bf4 100644 --- a/demo/download_models.py +++ b/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) diff --git a/demo/evaluation.py b/demo/evaluation.py index 606b798ccbfb467cead506ef55ff90983ea8c6ea..e5632f74e5ccf835cbee3d53cb858abeb2529793 100644 --- a/demo/evaluation.py +++ b/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") diff --git a/demo/indexing.py b/demo/indexing.py index c52b5da5b73cace54b97e2cab83a699d228ba0b6..0c3d3433eddb8f0592758771d83691d831349056 100644 --- a/demo/indexing.py +++ b/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], diff --git a/demo/qdrant_utils.py b/demo/qdrant_utils.py index 7581587e2fcbfdbab020f388f3bf95c50c3fb031..59169abd658b8b1df51ff1529b3589116041e558 100644 --- a/demo/qdrant_utils.py +++ b/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()}" diff --git a/demo/test_qdrant_connection.py b/demo/test_qdrant_connection.py index dabae1fc35f89da9bdb41c5b89ab8995989c3e9b..34fb5e120dee58a99effe530c0af3f1c31a03c7a 100644 --- a/demo/test_qdrant_connection.py +++ b/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__": diff --git a/demo/ui/__init__.py b/demo/ui/__init__.py index eeedf788dd2ce32b1ba4b69befd17188377b91b5..4b0dd8b192afb6bf8ce4e9bb0c0c221879afc941 100644 --- a/demo/ui/__init__.py +++ b/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", diff --git a/demo/ui/benchmark.py b/demo/ui/benchmark.py index c6893d12535bffedb2a251acbcd8b71e66ac2f03..3c9a077a5b0f51201a56eb857151035dc94bdd47 100644 --- a/demo/ui/benchmark.py +++ b/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("


", 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("


", 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) diff --git a/demo/ui/header.py b/demo/ui/header.py index c9c2780b98ea06907299941bc84033364f00790c..59a646fe3807031e2081ec8977a6a4a26029df14 100644 --- a/demo/ui/header.py +++ b/demo/ui/header.py @@ -4,7 +4,8 @@ import streamlit as st def render_header(): - st.markdown(""" + st.markdown( + """

{model_short} model", unsafe_allow_html=True) - + st.markdown( + f"✅ Found {model_short} 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) diff --git a/demo/ui/sidebar.py b/demo/ui/sidebar.py index 47a6508e992cf31bb2907426aafa9caceef10dda..12eaf213d943147586098570f679c7c15200ee5b 100644 --- a/demo/ui/sidebar.py +++ b/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( + """ - """, 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"{vname}{dim_str}, {dtype}, {disk_icon}") + rows.append( + f"{vname}{dim_str}, {dtype}, {disk_icon}" + ) table_html = f"{''.join(rows)}
" 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() diff --git a/demo/ui/upload.py b/demo/ui/upload.py index 1ae0d19ce4500bebbc8741d05a7bf951e3b65bbf..03e3e78bbf96afd93415c1782d3afb686cc394f2 100644 --- a/demo/ui/upload.py +++ b/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: diff --git a/demo_app.py b/demo_app.py new file mode 100644 index 0000000000000000000000000000000000000000..2eb00c28d8d890f4ec5f1b3c8634e0afccfb6853 --- /dev/null +++ b/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( + """ +
+

+ 🔬 Visual RAG Toolkit +

+

+ SIGIR 2026 Demo - Multi-Vector Visual Document Retrieval +

+
+ """, + 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"{vname}{dim_str}, {dtype}, {disk_icon}" + ) + table_html = f"{''.join(rows)}
" + 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 {model_short} 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() diff --git a/examples/COMMANDS.md b/examples/COMMANDS.md index 9329fb0f05e11a4ec36bacbbe656bdd3d88d5e51..0451ae93accfce360225c888565bb34f4a2eb29a 100644 --- a/examples/COMMANDS.md +++ b/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 +``` + diff --git a/examples/process_pdfs.py b/examples/process_pdfs.py index 7da3e0a62b7a54438f040f9b02f6076efdbdc18c..483674c67b2235c0d99da0b4e71d468a06bec793 100644 --- a/examples/process_pdfs.py +++ b/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() - - - - - - - diff --git a/examples/search_demo.py b/examples/search_demo.py index bb486c3dfb81236e5a2942ca8d206ef186d397fc..38b21eef12098b2f9b019850e906f826a8368fb6 100644 --- a/examples/search_demo.py +++ b/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() - - - - - - - diff --git a/scripts/colqwen25_probe.py b/scripts/colqwen25_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..05214925b7562f7662ae2ab7920f246ba583c321 --- /dev/null +++ b/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() diff --git a/scripts/compare_eval_scopes.py b/scripts/compare_eval_scopes.py new file mode 100644 index 0000000000000000000000000000000000000000..b2a50ebcd061e4573ef53c779c0dfa48e0355a02 --- /dev/null +++ b/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() diff --git a/scripts/compare_models_sample_queries.py b/scripts/compare_models_sample_queries.py new file mode 100644 index 0000000000000000000000000000000000000000..4fb6d09d0bdb4456b86da2a7dd56b569cc1cee2d --- /dev/null +++ b/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() diff --git a/scripts/create_qdrant_payload_indexes.py b/scripts/create_qdrant_payload_indexes.py new file mode 100644 index 0000000000000000000000000000000000000000..43b60307b25047cb879810b5a138075b65cd69e4 --- /dev/null +++ b/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() diff --git a/scripts/debug_failed_docs.py b/scripts/debug_failed_docs.py new file mode 100644 index 0000000000000000000000000000000000000000..f787be09dedbdc8bed6bbf98c58738917099b4e0 --- /dev/null +++ b/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() diff --git a/scripts/debug_vidore_qrels_alignment.py b/scripts/debug_vidore_qrels_alignment.py new file mode 100644 index 0000000000000000000000000000000000000000..8d440bff5aeacf0254415b7c31cd302180865e7d --- /dev/null +++ b/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() diff --git a/scripts/dedupe_failure_logs.py b/scripts/dedupe_failure_logs.py new file mode 100644 index 0000000000000000000000000000000000000000..c96654477ba8dcface3fee7ec8f136856f59779b --- /dev/null +++ b/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() diff --git a/scripts/force_qdrant_reindex.py b/scripts/force_qdrant_reindex.py new file mode 100644 index 0000000000000000000000000000000000000000..7df68d40c8a9578e4d06de2ea31017cddc796622 --- /dev/null +++ b/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() diff --git a/scripts/inspect_qdrant_collection.py b/scripts/inspect_qdrant_collection.py new file mode 100644 index 0000000000000000000000000000000000000000..0125ffa3724a013c07455ce852459c513b9557d2 --- /dev/null +++ b/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() diff --git a/scripts/qdrant_clone_collection_no_index.py b/scripts/qdrant_clone_collection_no_index.py new file mode 100644 index 0000000000000000000000000000000000000000..4d1d7caa8ba6881aa3789ffedcfb261e7528bb73 --- /dev/null +++ b/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() diff --git a/scripts/qdrant_debug_collection.py b/scripts/qdrant_debug_collection.py new file mode 100644 index 0000000000000000000000000000000000000000..b7e3e1e30ecbbc63ecfe996e3959561a71eed992 --- /dev/null +++ b/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 \\ + --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 \\ + --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() diff --git a/scripts/qdrant_disable_hnsw.py b/scripts/qdrant_disable_hnsw.py new file mode 100644 index 0000000000000000000000000000000000000000..dddeaf6d5d9a3c9f0e42ab0adba49a426a860188 --- /dev/null +++ b/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() diff --git a/scripts/qdrant_modify_vectors_smoketest.py b/scripts/qdrant_modify_vectors_smoketest.py new file mode 100644 index 0000000000000000000000000000000000000000..fa47211f9b51527f9565c698822b4668a851fa6b --- /dev/null +++ b/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() diff --git a/scripts/qdrant_rebuild_collection_no_index.py b/scripts/qdrant_rebuild_collection_no_index.py new file mode 100644 index 0000000000000000000000000000000000000000..e15d6768129c5f6a687de23ad963912844f0b953 --- /dev/null +++ b/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() diff --git a/scripts/qdrant_recompute_colqwen_pooling_from_initial.py b/scripts/qdrant_recompute_colqwen_pooling_from_initial.py new file mode 100644 index 0000000000000000000000000000000000000000..6e1f947abef34e3ef446e11f82f923d4fb1a08d0 --- /dev/null +++ b/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() diff --git a/scripts/query_token_stats.py b/scripts/query_token_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..cbc713201c9e1c6d8b56e6ef7a1351e479cd498a --- /dev/null +++ b/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() diff --git a/scripts/update_qdrant_indexing_threshold.py b/scripts/update_qdrant_indexing_threshold.py new file mode 100644 index 0000000000000000000000000000000000000000..0f40d322c4e9de1670b0be8b4452ce996cf9c81d --- /dev/null +++ b/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() diff --git a/tests/__init__.py b/tests/__init__.py index befc3ec915ed91e76f563ba6f4eba0a41d21b897..461f11f29b97d12771a884ea4450e8e45996428d 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,8 +1 @@ # Tests for visual-rag-toolkit - - - - - - - diff --git a/tests/test_config.py b/tests/test_config.py index a2ae9698740598b02a701ed83e65b9dce176d8f0..80bf2eace2df27a7bd2c44239ad407d0bc13455b 100644 --- a/tests/test_config.py +++ b/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 == {} - - - - - - - diff --git a/tests/test_pdf_processor.py b/tests/test_pdf_processor.py index 84c7f4eee5b85c57ba50aa1ac3281fdedc035a68..ef8cd542b7bae6cc5f86172aaef09542c8e9f6d3 100644 --- a/tests/test_pdf_processor.py +++ b/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}") - - - - - - - diff --git a/tests/test_pooling.py b/tests/test_pooling.py index 0b66b5d02f511487870d38079ba6fc3a872cf18f..8a77a88c8217a542d278e27d94d46bd4e432b042 100644 --- a/tests/test_pooling.py +++ b/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) diff --git a/tests/test_strategies.py b/tests/test_strategies.py index 100f197dd727224a328ef1237b159ee2165c9942..c019b3418ac7d0842ce87d95631b9dcdd6f2e91f 100644 --- a/tests/test_strategies.py +++ b/tests/test_strategies.py @@ -1,56 +1,56 @@ """Tests for embedding strategies (pooling vs standard).""" -import pytest import numpy as np +import pytest class TestEmbeddingStrategies: """Test different embedding strategies.""" - + def test_pooling_strategy_reduces_size(self): """Pooling strategy should produce smaller embeddings.""" # Simulating what pipeline does for pooling strategy full_embedding = np.random.randn(1000, 128).astype(np.float32) visual_indices = list(range(0, 832)) # Visual tokens only - + # Pooling: extract visual tokens visual_embedding = full_embedding[visual_indices] - + assert visual_embedding.shape[0] < full_embedding.shape[0] assert visual_embedding.shape[0] == 832 - + def test_standard_strategy_preserves_all(self): """Standard strategy should keep all tokens.""" full_embedding = np.random.randn(1000, 128).astype(np.float32) - + # Standard: use all tokens embedding_for_storage = full_embedding - + assert embedding_for_storage.shape == full_embedding.shape - + def test_pooling_produces_tile_vectors(self): """Pooling should produce [num_tiles, dim] vectors.""" from visual_rag.embedding.pooling import tile_level_mean_pooling - + # Visual tokens: 13 tiles × 64 patches visual_embedding = np.random.randn(832, 128).astype(np.float32) - + pooled = tile_level_mean_pooling(visual_embedding, num_tiles=13, patches_per_tile=64) - + # Should be [num_tiles, 128] assert pooled.shape == (13, 128) - + def test_standard_produces_single_vector(self): """Standard pooling should produce [1, dim] for mean_pooling vector.""" from visual_rag.embedding.pooling import global_mean_pooling - + full_embedding = np.random.randn(1000, 128).astype(np.float32) - + pooled = global_mean_pooling(full_embedding) - + # Should be [128] - can reshape to [1, 128] assert pooled.shape == (128,) - + # For storage as multi-vector pooled_for_storage = pooled.reshape(1, -1) assert pooled_for_storage.shape == (1, 128) @@ -58,22 +58,22 @@ class TestEmbeddingStrategies: class TestStrategyValidation: """Test strategy validation in pipeline.""" - + def test_valid_strategies(self): """Valid strategies should be accepted.""" from visual_rag.indexing.pipeline import ProcessingPipeline - + assert "pooling" in ProcessingPipeline.STRATEGIES assert "standard" in ProcessingPipeline.STRATEGIES assert "all" in ProcessingPipeline.STRATEGIES - + def test_invalid_strategy_raises(self): """Invalid strategy should raise ValueError.""" from visual_rag.indexing.pipeline import ProcessingPipeline - + with pytest.raises(ValueError, match="Invalid embedding_strategy"): ProcessingPipeline(embedding_strategy="invalid_strategy") - + def test_all_strategy_avoids_double_embedding(self): """'all' strategy computes embedding once, stores both representations.""" # The 'all' strategy is efficient because: @@ -87,15 +87,14 @@ class TestStrategyValidation: class TestStrategyMetadata: """Test that strategy is stored in metadata.""" - + def test_metadata_contains_strategy(self): """Metadata should include which strategy was used.""" # This is important for paper comparison - knowing which # strategy produced which results - - expected_fields = ["embedding_strategy", "num_visual_tokens", "total_tokens"] - + + _expected_fields = ["embedding_strategy", "num_visual_tokens", "total_tokens"] + # These fields should be in the metadata stored in Qdrant # Verified via pipeline code review pass # Marker test - actual verification is integration test - diff --git a/visual_rag/cli/main.py b/visual_rag/cli/main.py index 0c57237872024b830254fd51b9bbc2d990260948..98d6def97d9751836ff96afedbde144a811969aa 100644 --- a/visual_rag/cli/main.py +++ b/visual_rag/cli/main.py @@ -561,8 +561,19 @@ Examples: search_parser.add_argument( "--stage1-mode", type=str, - default="pooled_query_vs_tiles", - choices=["pooled_query_vs_tiles", "tokens_vs_tiles", "pooled_query_vs_global"], + default="pooled_query_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 (deprecated) + "pooled_query_vs_tiles", + "tokens_vs_tiles", + "pooled_query_vs_experimental", + "tokens_vs_experimental", + ], help="Stage 1 mode for two-stage retrieval", ) search_parser.add_argument("--year", type=int, help="Filter by year") diff --git a/visual_rag/demo_runner.py b/visual_rag/demo_runner.py index e54394e506a5eb890a5a51618fe22b0c69ecc6e2..12bf12d64a28e73ed0f5383f047d74bdbdd236a5 100644 --- a/visual_rag/demo_runner.py +++ b/visual_rag/demo_runner.py @@ -65,6 +65,7 @@ def demo( # Make sure the demo doesn't spam internal Streamlit warnings in logs. env.setdefault("STREAMLIT_BROWSER_GATHER_USAGE_STATS", "false") + print("Launching Streamlit demo:", " ".join(cmd), file=sys.stderr, flush=True) return subprocess.call(cmd, env=env) @@ -86,3 +87,7 @@ def main() -> None: extra_args=unknown, ) raise SystemExit(rc) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/visual_rag/embedding/pooling.py b/visual_rag/embedding/pooling.py index c3a52e61e5babb118eece199c94bc7edda7fa91d..f9a5b10912110cab216ca0af68c5d8480821df6c 100644 --- a/visual_rag/embedding/pooling.py +++ b/visual_rag/embedding/pooling.py @@ -124,6 +124,67 @@ def colpali_row_mean_pooling( return pooled.astype(out_dtype) +def adaptive_row_mean_pooling_from_grid( + embedding: Union[torch.Tensor, np.ndarray], + *, + grid_h: int, + grid_w: int, + target_rows: int = 32, + output_dtype: Optional[np.dtype] = None, +) -> np.ndarray: + """ + Row-mean pooling for arbitrary H×W patch grids with adaptive down/up-sampling to `target_rows`. + + This is useful for dynamic-resolution models (e.g., ColQwen2.5) where the number of + visual tokens (patches) is not fixed to a 32×32 grid. + + Steps: + 1) reshape tokens to [H, W, dim] + 2) mean over columns -> [H, dim] + 3) adaptive mean-pool rows to `target_rows` -> [target_rows, dim] + """ + out_dtype = _infer_output_dtype(embedding, output_dtype) + if isinstance(embedding, torch.Tensor): + if embedding.dtype == torch.bfloat16: + emb_np = embedding.cpu().float().numpy() + else: + emb_np = embedding.cpu().numpy().astype(np.float32) + else: + emb_np = np.array(embedding, dtype=np.float32) + + num_tokens, dim = emb_np.shape + expected = int(grid_h) * int(grid_w) + if num_tokens != expected: + raise ValueError( + f"Expected {expected} visual tokens for grid_h×grid_w={grid_h}×{grid_w}, got {num_tokens}" + ) + + grid = emb_np.reshape(int(grid_h), int(grid_w), int(dim)) + rows = grid.mean(axis=1) # [H, dim] + + h = int(rows.shape[0]) + target_rows = int(target_rows) + if target_rows <= 0: + raise ValueError("target_rows must be > 0") + if h == target_rows: + return rows.astype(out_dtype) + if h == 1: + return np.repeat(rows, repeats=target_rows, axis=0).astype(out_dtype) + + # Adaptive average pooling along the row dimension. + # We use evenly spaced bins over [0, H) and mean rows per bin. + edges = np.linspace(0, h, target_rows + 1) + pooled = np.zeros((target_rows, int(dim)), dtype=np.float32) + for i in range(target_rows): + start = int(np.floor(edges[i])) + end = int(np.ceil(edges[i + 1])) + start = max(0, min(start, h - 1)) + end = max(start + 1, min(end, h)) + pooled[i] = rows[start:end].mean(axis=0) + + return pooled.astype(out_dtype) + + def colsmol_experimental_pooling( embedding: Union[torch.Tensor, np.ndarray], num_tiles: int, @@ -173,20 +234,19 @@ def colsmol_experimental_pooling( def colpali_experimental_pooling_from_rows( row_vectors: Union[torch.Tensor, np.ndarray], + *, + window_size: int = 3, output_dtype: Optional[np.dtype] = None, ) -> np.ndarray: """ - Experimental "convolution-style" pooling with window size 3. + Experimental "convolution-style" pooling with an odd `window_size` (default: 3). - For N input rows, produces N + 2 output vectors: - - Position 0: row[0] alone (1 row) - - Position 1: mean(rows[0:2]) (2 rows) - - Position 2: mean(rows[0:3]) (3 rows) - - Positions 3 to N-1: sliding window of 3 (rows[i-2:i+1]) - - Position N: mean(rows[N-2:N]) (last 2 rows) - - Position N+1: row[N-1] alone (last row) + For N input rows and radius r = window_size//2, produces N + 2r output vectors. + Each output position uses a clipped window around a (possibly out-of-range) center: + center = i - r, i in [0, N + 2r - 1] + window = rows[max(0, center-r) : min(N, center+r+1)] - For N=32 rows: produces 34 vectors. + For window_size=3 and N=32 rows: produces 34 vectors (same as previous implementation). """ out_dtype = _infer_output_dtype(row_vectors, output_dtype) if isinstance(row_vectors, torch.Tensor): @@ -200,30 +260,29 @@ def colpali_experimental_pooling_from_rows( n, dim = rows.shape if n < 1: raise ValueError("row_vectors must be non-empty") - if n == 1: + + window_size = int(window_size) + if window_size < 1: + raise ValueError("window_size must be >= 1") + if window_size % 2 == 0: + raise ValueError("window_size must be odd") + if window_size == 1 or n == 1: return rows.astype(out_dtype) - if n == 2: - return np.stack([rows[0], rows[:2].mean(axis=0), rows[1]], axis=0).astype(out_dtype) - if n == 3: - return np.stack( - [ - rows[0], - rows[:2].mean(axis=0), - rows[:3].mean(axis=0), - rows[1:3].mean(axis=0), - rows[2], - ], - axis=0, - ).astype(out_dtype) - - out = np.zeros((n + 2, dim), dtype=np.float32) - out[0] = rows[0] - out[1] = rows[:2].mean(axis=0) - out[2] = rows[:3].mean(axis=0) - for i in range(3, n): - out[i] = rows[i - 2 : i + 1].mean(axis=0) - out[n] = rows[n - 2 : n].mean(axis=0) - out[n + 1] = rows[n - 1] + + r = window_size // 2 + # Preserve legacy edge-case behavior for the default window_size=3: + # - n=1 -> (1, dim) + # - n=2 -> (3, dim): [row0, mean(row0,row1), row1] + # - n>=3 -> (n+2, dim) + if int(window_size) == 3 and int(n) == 2: + mid = rows.mean(axis=0) + return np.stack([rows[0], mid, rows[1]], axis=0).astype(out_dtype) + out = np.zeros((n + 2 * r, dim), dtype=np.float32) + for i in range(n + 2 * r): + center = i - r + lo = max(0, center - r) + hi = min(n - 1, center + r) + out[i] = rows[lo : hi + 1].mean(axis=0) return out.astype(out_dtype) diff --git a/visual_rag/embedding/visual_embedder.py b/visual_rag/embedding/visual_embedder.py index addd3db9f4a00d9a9f029f35fd4c39bb8f1555a8..1198f8684ad1c2f9dbf74fd887836447d7d6b026 100644 --- a/visual_rag/embedding/visual_embedder.py +++ b/visual_rag/embedding/visual_embedder.py @@ -12,8 +12,11 @@ The embedder is BACKEND-AGNOSTIC - configure which model to use via the """ import gc +import json import logging import os +from datetime import datetime, timezone +from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union import numpy as np @@ -164,6 +167,12 @@ class VisualEmbedder: "pip install visual-rag-toolkit[embedding] or " "pip install colpali-engine" ) + try: + # Newer colpali-engine versions add ColQwen2.5 support + from colpali_engine.models import ColQwen2_5, ColQwen2_5_Processor + except Exception: + ColQwen2_5 = None + ColQwen2_5_Processor = None logger.info(f"🤖 Loading ColPali model: {self.model_name}") logger.info(f" Device: {self.device}, dtype: {self.torch_dtype}") @@ -201,7 +210,54 @@ class VisualEmbedder: logger.info("✅ Loaded ColPali backend") return - if model_type.startswith("qwen2") or "colqwen" in (self.model_name or "").lower(): + model_lower = (self.model_name or "").lower() + is_qwen25 = ( + "colqwen2.5" in model_lower + or "colqwen2_5" in model_lower + or "qwen2_5" in model_type + or "qwen2.5" in model_type + ) + if is_qwen25: + if ColQwen2_5 is None or ColQwen2_5_Processor is None: + raise ImportError( + "ColQwen2.5 requires a newer colpali-engine. Install/upgrade with:\n" + ' pip install "transformers>=4.45.0"\n' + " pip install git+https://github.com/illuin-tech/colpali\n" + "or ensure colpali-engine>=0.3.7 is installed." + ) + attn_implementation = None + if self.device != "cpu": + try: + from transformers.utils.import_utils import is_flash_attn_2_available + + if is_flash_attn_2_available(): + attn_implementation = "flash_attention_2" + except Exception: + pass + self._model = ColQwen2_5.from_pretrained( + self.model_name, + torch_dtype=self.torch_dtype, + device_map=self.device, + attn_implementation=attn_implementation, + ).eval() + try: + self._processor = ColQwen2_5_Processor.from_pretrained( + self.model_name, **_processor_kwargs() + ) + except TypeError: + self._processor = ColQwen2_5_Processor.from_pretrained(self.model_name) + except Exception: + if self.processor_speed == "fast": + self._processor = ColQwen2_5_Processor.from_pretrained( + self.model_name, use_fast=False + ) + else: + raise + self._image_token_id = self._processor.image_token_id + logger.info("✅ Loaded ColQwen2.5 backend") + return + + if model_type.startswith("qwen2") or "colqwen" in model_lower: self._model = ColQwen2.from_pretrained( self.model_name, dtype=self.torch_dtype, @@ -350,6 +406,9 @@ class VisualEmbedder: batch_size = batch_size or self.batch_size outputs: List[torch.Tensor] = [] + fallback_count = 0 + nan_log_path: Optional[Path] = None + nan_logged = 0 iterator = range(0, len(query_texts), batch_size) if show_progress: iterator = tqdm(iterator, desc="📝 Embedding queries", unit="batch") @@ -374,6 +433,63 @@ class VisualEmbedder: non_special_mask = ids >= 4 if non_special_mask.any(): emb = emb[non_special_mask] + # ColQwen2.5 on MPS can produce NaNs when batching queries. + # If we detect NaNs/Infs, recompute the query embedding individually (stable). + try: + has_nan = bool(torch.isnan(emb).any()) + has_inf = bool(torch.isinf(emb).any()) + if has_nan or has_inf: + # Persist a reproducible sample for debugging. + try: + if nan_log_path is None: + log_dir = os.getenv("VISUALRAG_NAN_LOG_DIR") or str( + Path("results") / "nan_samples" + ) + Path(log_dir).mkdir(parents=True, exist_ok=True) + safe_model = ( + str(self.model_name or "model") + .replace("/", "_") + .replace(" ", "_") + .replace(":", "_") + ) + ts = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + nan_log_path = ( + Path(log_dir) / f"nan_queries__{safe_model}__{ts}.jsonl" + ) + + rec = { + "ts": datetime.now(timezone.utc).isoformat(), + "model_name": str(self.model_name), + "device": str(self.device), + "torch_dtype": str(self.torch_dtype), + "output_dtype": str(self.output_dtype), + "processor_speed": str( + getattr(self, "processor_speed", "unknown") + ), + "filter_special_tokens": bool(should_filter), + "batch_size": int(batch_size), + "global_query_index": int(i + j), + "query_text": str(batch[j]), + "has_nan": bool(has_nan), + "has_inf": bool(has_inf), + "torch_version": str(getattr(torch, "__version__", "")), + } + with nan_log_path.open("a", encoding="utf-8") as f: + f.write(json.dumps(rec, ensure_ascii=False) + "\n") + nan_logged += 1 + if nan_logged <= 3: + logger.warning( + "NaN/Inf detected in batched query embedding (idx=%d). " + "Logged sample to %s. Recomputing this query individually.", + int(i + j), + str(nan_log_path), + ) + except Exception: + pass + fallback_count += 1 + emb = self.embed_query(batch[j], filter_special_tokens=should_filter) + except Exception: + pass outputs.append(emb) else: outputs.extend(batch_embeddings) @@ -385,6 +501,17 @@ class VisualEmbedder: elif torch.backends.mps.is_available(): torch.mps.empty_cache() + if fallback_count > 0: + logger.warning( + "embed_queries(): detected NaN/Inf in %d/%d queries; " + "recomputed those queries individually for stability.", + int(fallback_count), + int(len(query_texts)), + ) + if nan_log_path is not None and nan_logged > 0: + logger.warning( + "NaN/Inf samples written to %s (%d rows).", str(nan_log_path), int(nan_logged) + ) return outputs def embed_images( @@ -441,6 +568,12 @@ class VisualEmbedder: input_ids = processed["input_ids"] batch_n_rows = processed.get("n_rows") batch_n_cols = processed.get("n_cols") + # Qwen2/2.5-VL style grid information (T, H, W) + batch_grid_thw = processed.get("image_grid_thw", None) + if batch_grid_thw is None: + batch_grid_thw = processed.get("grid_thw", None) + if batch_grid_thw is None: + batch_grid_thw = processed.get("image_grid", None) for j in range(input_ids.shape[0]): # Find visual token indices @@ -449,6 +582,30 @@ class VisualEmbedder: n_rows = batch_n_rows[j].item() if batch_n_rows is not None else None n_cols = batch_n_cols[j].item() if batch_n_cols is not None else None + grid_t = grid_h = grid_w = None + grid_h_eff = grid_w_eff = None + if batch_grid_thw is not None: + try: + g = batch_grid_thw[j] + # sometimes [B, N, 3] -> take first image + if hasattr(g, "dim") and g.dim() == 2: + g = g[0] + t, h, w = [int(x) for x in g.detach().cpu().tolist()] + grid_t, grid_h, grid_w = t, h, w + # ColQwen2.5/Qwen2.5-VL uses a 2×2 spatial merge internally, but different + # processor versions expose different grids: + # - Some expose the *post-merge* token grid (H×W == num_visual_tokens) + # - Others expose the *pre-merge* pixel/patch grid ((H/2)×(W/2) == num_visual_tokens) + # We infer the effective grid by matching the observed token count. + num_visual = int(len(visual_indices)) + if int(h) * int(w) == num_visual: + grid_h_eff, grid_w_eff = int(h), int(w) + elif ( + h % 2 == 0 and w % 2 == 0 and (h // 2) * (w // 2) == num_visual + ): + grid_h_eff, grid_w_eff = int(h // 2), int(w // 2) + except Exception: + pass token_infos.append( { @@ -457,6 +614,11 @@ class VisualEmbedder: "n_rows": n_rows, "n_cols": n_cols, "num_tiles": (n_rows * n_cols + 1) if n_rows and n_cols else None, + "grid_t": grid_t, + "grid_h": grid_h, + "grid_w": grid_w, + "grid_h_eff": grid_h_eff, + "grid_w_eff": grid_w_eff, } ) @@ -516,12 +678,29 @@ class VisualEmbedder: visual_embedding: Union[torch.Tensor, np.ndarray], token_info: Optional[Dict[str, Any]] = None, *, - target_vectors: int = 32, + target_vectors: Optional[int] = 32, ) -> np.ndarray: - from visual_rag.embedding.pooling import colpali_row_mean_pooling, tile_level_mean_pooling + from visual_rag.embedding.pooling import ( + adaptive_row_mean_pooling_from_grid, + colpali_row_mean_pooling, + tile_level_mean_pooling, + ) model_lower = (self.model_name or "").lower() is_colsmol = "colsmol" in model_lower + is_colqwen25 = "colqwen2.5" in model_lower or "colqwen2_5" in model_lower + target_vectors_cap: Optional[int] + if target_vectors is None: + target_vectors_cap = None + else: + try: + tv = int(target_vectors) + except Exception: + tv = 32 + target_vectors_cap = None if tv <= 0 else tv + # For non-dynamic models, default to the historical fixed 32 vectors when unset. + if not is_colqwen25 and target_vectors_cap is None: + target_vectors_cap = 32 if isinstance(visual_embedding, torch.Tensor): if visual_embedding.dtype == torch.bfloat16: @@ -539,19 +718,60 @@ class VisualEmbedder: visual_np, num_tiles=num_tiles, patches_per_tile=64, output_dtype=self.output_dtype ) + # ColQwen2.5 supports dynamic resolutions. The processor provides a pre-merge grid (grid_h/grid_w), + # but the *effective* token grid is (grid_h_eff, grid_w_eff) due to 2×2 spatial merge. + # We follow the dynamic shape by default: + # - if target_vectors is unset (None or <=0), return all effective rows (no cap, no upsampling) + # - else, return <= target_vectors rows (no upsampling). num_tokens = int(visual_np.shape[0]) + if is_colqwen25: + grid_h_eff = (token_info or {}).get("grid_h_eff") + grid_w_eff = (token_info or {}).get("grid_w_eff") + if grid_h_eff and grid_w_eff and int(grid_h_eff) * int(grid_w_eff) == int(num_tokens): + # Compute row means over the *effective* grid. + target_rows = int(grid_h_eff) + if target_vectors_cap is not None: + target_rows = min(int(target_vectors_cap), int(grid_h_eff)) + pooled_rows = adaptive_row_mean_pooling_from_grid( + visual_np, + grid_h=int(grid_h_eff), + grid_w=int(grid_w_eff), + target_rows=target_rows, + output_dtype=self.output_dtype, + ) + return pooled_rows + + # Fallback: infer a square grid if possible grid = int(round(float(num_tokens) ** 0.5)) - if grid * grid != num_tokens: - raise ValueError( - f"Cannot infer square grid from num_visual_tokens={num_tokens} for model={self.model_name}" + if grid * grid == num_tokens: + # For ColQwen2.5 with unset cap, keep all rows (grid) rather than defaulting to 32. + effective_target_rows = ( + int(grid) if (is_colqwen25 and target_vectors_cap is None) else int(target_vectors_cap) ) - if int(target_vectors) != int(grid): - raise ValueError( - f"target_vectors={target_vectors} does not match inferred grid_size={grid} for model={self.model_name}" + if int(grid) == int(effective_target_rows): + return colpali_row_mean_pooling( + visual_np, grid_size=int(effective_target_rows), output_dtype=self.output_dtype + ) + return adaptive_row_mean_pooling_from_grid( + visual_np, + grid_h=int(grid), + grid_w=int(grid), + target_rows=int(effective_target_rows), + output_dtype=self.output_dtype, ) - return colpali_row_mean_pooling( - visual_np, grid_size=int(target_vectors), output_dtype=self.output_dtype - ) + + # Last-resort: treat tokens as a sequence and adaptively mean-pool chunks to target_vectors rows. + # If unset (None/<=0), fall back to 32 to avoid returning extremely large multi-vectors. + tv_last = int(target_vectors_cap or 32) + edges = np.linspace(0, num_tokens, tv_last + 1) + pooled = np.zeros((tv_last, int(visual_np.shape[1])), dtype=np.float32) + for i in range(tv_last): + s = int(np.floor(edges[i])) + e = int(np.ceil(edges[i + 1])) + s = max(0, min(s, num_tokens - 1)) + e = max(s + 1, min(e, num_tokens)) + pooled[i] = visual_np[s:e].mean(axis=0) + return pooled.astype(self.output_dtype) def global_pool_from_mean_pool(self, mean_pool: np.ndarray) -> np.ndarray: if mean_pool.size == 0: @@ -563,7 +783,7 @@ class VisualEmbedder: visual_embedding: Union[torch.Tensor, np.ndarray], token_info: Optional[Dict[str, Any]] = None, *, - target_vectors: int = 32, + target_vectors: Optional[int] = 32, mean_pool: Optional[np.ndarray] = None, ) -> np.ndarray: from visual_rag.embedding.pooling import ( @@ -573,6 +793,7 @@ class VisualEmbedder: model_lower = (self.model_name or "").lower() is_colsmol = "colsmol" in model_lower + is_colqwen25 = "colqwen2.5" in model_lower or "colqwen2_5" in model_lower if isinstance(visual_embedding, torch.Tensor): if visual_embedding.dtype == torch.bfloat16: @@ -611,11 +832,11 @@ class VisualEmbedder: visual_np, token_info, target_vectors=target_vectors ) ) - if int(rows.shape[0]) != int(target_vectors): - raise ValueError( - f"experimental pooling expects mean_pool to have {target_vectors} rows, got {rows.shape[0]} for model={self.model_name}" - ) - return colpali_experimental_pooling_from_rows(rows, output_dtype=self.output_dtype) + # For ColPali we usually expect fixed 32 rows; for ColQwen2.5 the row count is dynamic (<= target_vectors). + window = 5 if is_colqwen25 else 3 + return colpali_experimental_pooling_from_rows( + rows, window_size=window, output_dtype=self.output_dtype + ) # Backward compatibility alias diff --git a/visual_rag/indexing/qdrant_indexer.py b/visual_rag/indexing/qdrant_indexer.py index a01dbdaae04122c7b019f3810d0b2ccd40601bcc..8e8d67205c4e3cf01217fc9e4316bc29486678ec 100644 --- a/visual_rag/indexing/qdrant_indexer.py +++ b/visual_rag/indexing/qdrant_indexer.py @@ -208,6 +208,9 @@ class QdrantIndexer: self.client.create_collection( collection_name=self.collection_name, vectors_config=vectors_config, + optimizers_config=qdrant_models.OptimizersConfigDiff( + indexing_threshold=int(indexing_threshold), + ), ) # Create required payload index for skip_existing functionality @@ -246,11 +249,46 @@ class QdrantIndexer: if not fields: return + # Cache between calls so multi-dataset runs don't spam logs. + if not hasattr(self, "_ensured_payload_indexes"): + self._ensured_payload_indexes = set() + self._payload_indexes_skip_logged = False + + requested_fields: List[str] = [] + for fc in fields: + try: + requested_fields.append(str(fc["field"])) + except Exception: + continue + requested_set = set(requested_fields) + + # Qdrant exposes indexed payload fields in collection_info.payload_schema + existing_indexed: set[str] = set() + try: + info = self.client.get_collection(self.collection_name) + payload_schema = getattr(info, "payload_schema", None) or {} + if isinstance(payload_schema, dict): + existing_indexed = set(str(k) for k in payload_schema.keys()) + except Exception as e: + logger.debug(f"Could not read existing payload schema: {e}") + + already = existing_indexed | set(self._ensured_payload_indexes) + missing = requested_set - already + + if not missing: + # Log this only once per process to avoid repetition across datasets. + if not getattr(self, "_payload_indexes_skip_logged", False): + logger.info("📇 Payload indexes already exist — skipping creation") + self._payload_indexes_skip_logged = True + self._ensured_payload_indexes |= requested_set + return - logger.info("📇 Creating payload indexes...") + logger.info(f"📇 Creating payload indexes ({len(missing)} new)...") for field_config in fields: - field_name = field_config["field"] + field_name = str(field_config["field"]) + if field_name not in missing: + continue field_type_str = field_config.get("type", "keyword") field_type = type_mapping.get(field_type_str, qdrant_models.PayloadSchemaType.KEYWORD) @@ -260,8 +298,11 @@ class QdrantIndexer: field_name=field_name, field_schema=field_type, ) + self._ensured_payload_indexes.add(field_name) logger.info(f" ✅ {field_name} ({field_type_str})") except Exception as e: + # If Qdrant reports it already exists anyway, treat it as ensured. + self._ensured_payload_indexes.add(field_name) logger.debug(f" Index {field_name} might already exist: {e}") def upload_batch( @@ -383,8 +424,6 @@ class QdrantIndexer: wait=wait, ) - logger.info(f" ✅ Uploaded {len(points)} points to Qdrant") - if delay_between_batches > 0: if _is_cancelled(): return 0 diff --git a/visual_rag/qdrant_admin.py b/visual_rag/qdrant_admin.py index 7543509610bc222356fc6caf686a042718d6e99e..84b92971f1c79721ea3bff94c09a77edb9cfb4a3 100644 --- a/visual_rag/qdrant_admin.py +++ b/visual_rag/qdrant_admin.py @@ -75,14 +75,25 @@ class QdrantAdmin: conn = _resolve_qdrant_connection(url=url, api_key=api_key) grpc_port = _infer_grpc_port(conn.url) if prefer_grpc else None - self.client = QdrantClient( - url=conn.url, - api_key=conn.api_key, - prefer_grpc=bool(prefer_grpc), - grpc_port=grpc_port, - timeout=int(timeout), - check_compatibility=False, - ) + + def _make(use_grpc: bool): + return QdrantClient( + url=conn.url, + api_key=conn.api_key, + prefer_grpc=bool(use_grpc), + grpc_port=grpc_port if use_grpc else None, + timeout=int(timeout), + check_compatibility=False, + ) + + self.client = _make(bool(prefer_grpc)) + if prefer_grpc: + # gRPC can fail in some environments (DNS, proxies, etc.). + # Fall back to REST for admin operations. + try: + _ = self.client.get_collections() + except Exception: + self.client = _make(False) def get_collection_info(self, *, collection_name: str) -> Dict[str, Any]: info = self.client.get_collection(collection_name) @@ -220,3 +231,44 @@ class QdrantAdmin: ) return self.get_collection_info(collection_name=collection_name) + + def ensure_collection_all_in_ram( + self, + *, + collection_name: str, + timeout: Optional[int] = None, + ) -> Dict[str, Any]: + """ + Best-effort configuration to keep vectors/indexes in RAM. + + Ensures: + - All existing named vectors have on_disk=False and hnsw_config.on_disk=False + - Collection hnsw_config.on_disk=False + - Collection params.on_disk_payload=False + + Note: This is configuration-level. Actual residency still depends on available RAM + and the OS page cache; Qdrant doesn't expose a "pin all vectors in RAM now" API. + """ + collection_name = str(collection_name) + info = self.client.get_collection(collection_name) + vectors = {} + try: + existing = list((info.config.params.vectors or {}).keys()) + except Exception: + existing = [] + for vname in existing: + vectors[str(vname)] = {"on_disk": False, "hnsw_config": {"on_disk": False}} + + if vectors: + self.modify_collection_vector_config( + collection_name=collection_name, vectors=vectors, timeout=timeout + ) + + self.modify_collection_config( + collection_name=collection_name, + hnsw_config={"on_disk": False}, + collection_params={"on_disk_payload": False}, + timeout=timeout, + ) + + return self.get_collection_info(collection_name=collection_name) diff --git a/visual_rag/retrieval/multi_vector.py b/visual_rag/retrieval/multi_vector.py index b0b3945c9fbe2462a786b36637364aac1bdbfa4d..de52068cda728fed3584048a054ebabc127a6a5a 100644 --- a/visual_rag/retrieval/multi_vector.py +++ b/visual_rag/retrieval/multi_vector.py @@ -151,7 +151,7 @@ class MultiVectorRetriever: top_k: int = 10, mode: str = "single_full", prefetch_k: Optional[int] = None, - stage1_mode: str = "pooled_query_vs_tiles", + stage1_mode: str = "pooled_query_vs_standard_pooling", filter_obj=None, return_embeddings: bool = False, ) -> List[Dict[str, Any]]: @@ -179,7 +179,7 @@ class MultiVectorRetriever: top_k: int = 10, mode: str = "single_full", prefetch_k: Optional[int] = None, - stage1_mode: str = "pooled_query_vs_tiles", + stage1_mode: str = "pooled_query_vs_standard_pooling", stage1_k: Optional[int] = None, stage2_k: Optional[int] = None, filter_obj=None, @@ -190,14 +190,14 @@ class MultiVectorRetriever: query_embedding=query_embedding, top_k=top_k, filter_obj=filter_obj, - using="initial", + strategy="multi_vector", ) elif mode == "single_pooled": return self._single_stage.search( query_embedding=query_embedding, top_k=top_k, filter_obj=filter_obj, - using="mean_pooling", + strategy="pooled_tile", ) elif mode == "two_stage": return self._two_stage.search_server_side( diff --git a/visual_rag/retrieval/two_stage.py b/visual_rag/retrieval/two_stage.py index 07186448661e2852ff1e5f20dfe530fece340177..820664c0715a173bbfa0016c6dd8f7c423c6becd 100644 --- a/visual_rag/retrieval/two_stage.py +++ b/visual_rag/retrieval/two_stage.py @@ -105,7 +105,7 @@ class TwoStageRetriever: top_k: int = 10, prefetch_k: Optional[int] = None, filter_obj=None, - stage1_mode: str = "pooled_query_vs_tiles", + stage1_mode: str = "pooled_query_vs_standard_pooling", ) -> List[Dict[str, Any]]: """ Two-stage retrieval using Qdrant's native prefetch (all server-side). @@ -128,16 +128,26 @@ class TwoStageRetriever: if prefetch_k is None: prefetch_k = max(100, top_k * 10) + # Backwards-compatible aliases: if stage1_mode == "pooled_query_vs_tiles": + stage1_mode = "pooled_query_vs_standard_pooling" + elif stage1_mode == "tokens_vs_tiles": + stage1_mode = "tokens_vs_standard_pooling" + elif stage1_mode == "pooled_query_vs_experimental": + stage1_mode = "pooled_query_vs_experimental_pooling" + elif stage1_mode == "tokens_vs_experimental": + stage1_mode = "tokens_vs_experimental_pooling" + + if stage1_mode == "pooled_query_vs_standard_pooling": prefetch_query = query_np.mean(axis=0).tolist() prefetch_using = self.pooled_vector_name - elif stage1_mode == "tokens_vs_tiles": + elif stage1_mode == "tokens_vs_standard_pooling": prefetch_query = query_np.tolist() prefetch_using = self.pooled_vector_name - elif stage1_mode == "pooled_query_vs_experimental": + elif stage1_mode == "pooled_query_vs_experimental_pooling": prefetch_query = query_np.mean(axis=0).tolist() prefetch_using = self.experimental_vector_name - elif stage1_mode == "tokens_vs_experimental": + elif stage1_mode == "tokens_vs_experimental_pooling": prefetch_query = query_np.tolist() prefetch_using = self.experimental_vector_name elif stage1_mode == "pooled_query_vs_global": @@ -188,7 +198,7 @@ class TwoStageRetriever: filter_obj=None, use_reranking: bool = True, return_embeddings: bool = False, - stage1_mode: str = "pooled_query_vs_tiles", + stage1_mode: str = "pooled_query_vs_standard_pooling", ) -> List[Dict[str, Any]]: """ Two-stage retrieval: prefetch with pooling, rerank with MaxSim. @@ -201,9 +211,11 @@ class TwoStageRetriever: use_reranking: Enable stage 2 reranking (default: True) return_embeddings: Include embeddings in results stage1_mode: - - "pooled_query_vs_tiles": pool query to 1×dim and search tile vectors (using="mean_pooling") - - "tokens_vs_tiles": search tile vectors with full query tokens (using="mean_pooling") - - "pooled_query_vs_global": pool query to 1×dim and search global pooled doc vectors (using="global_pooling") + - "pooled_query_vs_standard_pooling": pool query to 1×dim and search standard_pooling vectors (Qdrant named vector: "mean_pooling") + - "tokens_vs_standard_pooling": search standard_pooling vectors with full query tokens (Qdrant named vector: "mean_pooling") + - "pooled_query_vs_experimental_pooling": pool query to 1×dim and search experimental_pooling vectors (Qdrant named vector: "experimental_pooling") + - "tokens_vs_experimental_pooling": search experimental_pooling vectors with full query tokens (Qdrant named vector: "experimental_pooling") + - "pooled_query_vs_global": pool query to 1×dim and search global pooled doc vectors (Qdrant named vector: "global_pooling") Returns: List of results with scores and metadata: @@ -318,7 +330,7 @@ class TwoStageRetriever: query_np: np.ndarray, top_k: int, filter_obj=None, - stage1_mode: str = "pooled_query_vs_tiles", + stage1_mode: str = "pooled_query_vs_standard_pooling", ) -> List[Dict[str, Any]]: """Stage 1: Prefetch candidates.""" if stage1_mode == "pooled_query_vs_tiles":