feat(benchmark): support multi-dataset evaluation with configurable top-k list

- Evaluate multiple datasets in a single run (CIFAR10 + CIFAR100)
- Report Recall@K for a list of K values from one underlying search
- Save results to disk: summary.md, metrics.csv, per-dataset predictions.csv, confusion_matrix.csv
- Richer evaluation output: query_labels, topk_ids, topk_labels for downstream analysis
- Add --dataset and --top-k CLI overrides for benchmark command
- Update config schema: dataset→datasets, top_k→top_k_list

.justfile

@@ -25,8 +25,12 @@ ssh:
     -L 127.0.0.1:2718:172.30.0.2:2718 \
     {{ remote_ssh_target }}
 
-docker:
+docker cmd="":
-    docker exec -it {{ remote_docker_container }} bash
+    @if [ -z "{{ cmd }}" ]; then \
+        docker exec -it -w {{ env("REMOTE_WORKDIR") }} {{ remote_docker_container }} bash; \
+    else \
+        docker exec -i -w {{ env("REMOTE_WORKDIR") }} {{ remote_docker_container }} bash -lc {{ quote(cmd) }}; \
+    fi
 
 marimo +notebook:
     uv run marimo edit {{ notebook }} --host 0.0.0.0 --port 2718 --no-token
@@ -84,36 +88,36 @@ cam-test-retrieval-write-noise:
 
 # Prepare CIFAR10 hash artifact for CAM retrieval smoke benchmark
 cam-prepare-retrieval-cifar10 ROWS="512" QUERIES="128":
-    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt"
+    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar10 --num-rows {{ ROWS }} --max-queries {{ QUERIES }} --compressor-path outputs/hash_compressor.pt"
 
 # Prepare CIFAR100 hash artifact for CAM retrieval benchmark
 cam-prepare-retrieval-cifar100 ROWS="512" QUERIES="128":
-    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ROWS}} --max-queries {{QUERIES}} --compressor-path outputs/hash_compressor.pt"
+    just remote "python scripts/prepare_cam_retrieval_dataset.py --dataset cifar100 --num-rows {{ ROWS }} --max-queries {{ QUERIES }} --compressor-path outputs/hash_compressor.pt"
 
 # Run CAM retrieval benchmark on a prepared artifact without hardware noise
 cam-test-retrieval-artifact DATASET_PATH NUM_ROWS="4096":
-    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"
+    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{ NUM_ROWS }} WRITE_NOISE_EN=0 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"
 
 # Run CAM retrieval benchmark on a prepared artifact with write noise enabled (Phase 2: read noise removed)
 cam-test-retrieval-artifact-write-noise DATASET_PATH NUM_ROWS="4096":
-    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{NUM_ROWS}} WRITE_NOISE_EN=1 WRITE_NOISE_RATE_NUM=1 WRITE_NOISE_RATE_DEN=100 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"
+    just remote "make -C hw/sim clean && make -C hw/sim test-benchmark-retrieval TOPK_K=5 NUM_ROWS={{ NUM_ROWS }} WRITE_NOISE_EN=1 WRITE_NOISE_RATE_NUM=1 WRITE_NOISE_RATE_DEN=100 CAM_RETRIEVAL_DATASET={{ DATASET_PATH }}"
 
 # ── CAM retrieval benchmark noise sweep ────────────────────────────────────────
-
 # Run noise sweep on a prepared dataset (0%–100%, step 10%)
+
 # Usage: just cam-benchmark-retrieval-sweep DATASET=outputs/.../cifar10_hash512_rows512_queries128.npz NUM_ROWS=512
 cam-benchmark-retrieval-sweep DATASET NUM_ROWS="512":
-    just remote "python scripts/run_retrieval_noise_sweep.py --dataset {{DATASET}} --num-rows {{NUM_ROWS}} --output docs/cam_retrieval_noise_sweep.md"
+    just remote "python scripts/run_retrieval_noise_sweep.py --dataset {{ DATASET }} --num-rows {{ NUM_ROWS }} --output docs/cam_retrieval_noise_sweep.md"
 
 # Prepare CIFAR10 dataset + run full noise sweep (all-in-one)
 cam-benchmark-sweep-cifar10 ROWS="512" QUERIES="128":
-    just cam-prepare-retrieval-cifar10 {{ROWS}} {{QUERIES}}
+    just cam-prepare-retrieval-cifar10 {{ ROWS }} {{ QUERIES }}
-    just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar10_hash512_rows{{ROWS}}_queries{{QUERIES}}.npz --num-rows {{ROWS}} --output docs/cam_retrieval_noise_sweep_cifar10.md"
+    just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar10_hash512_rows{{ ROWS }}_queries{{ QUERIES }}.npz --num-rows {{ ROWS }} --output docs/cam_retrieval_noise_sweep_cifar10.md"
 
 # Prepare CIFAR100 dataset + run full noise sweep (all-in-one)
 cam-benchmark-sweep-cifar100 ROWS="512" QUERIES="128":
-    just cam-prepare-retrieval-cifar100 {{ROWS}} {{QUERIES}}
+    just cam-prepare-retrieval-cifar100 {{ ROWS }} {{ QUERIES }}
-    just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows{{ROWS}}_queries{{QUERIES}}.npz --num-rows {{ROWS}} --output docs/cam_retrieval_noise_sweep_cifar100.md"
+    just remote "python scripts/run_retrieval_noise_sweep.py --dataset outputs/cam_retrieval_benchmark/datasets/cifar100_hash512_rows{{ ROWS }}_queries{{ QUERIES }}.npz --num-rows {{ ROWS }} --output docs/cam_retrieval_noise_sweep_cifar100.md"
 
 # ── Remote ↔ local sync ────────────────────────────────────────────────────────
 

mini-nav/benchmarks/runner.py

@@ -1,5 +1,8 @@
 """Benchmark runner for executing evaluations."""
 
+import csv
+import json
+from datetime import datetime
 from pathlib import Path
 from typing import Any, Callable, cast
 
@@ -23,7 +26,7 @@ def _create_task(config: BenchmarkConfig, model_config: ModelConfig | None) -> A
     Returns:
         Benchmark task instance.
     """
-    task_kwargs: dict[str, Any] = {"top_k": config.task.top_k}
+    task_kwargs: dict[str, Any] = {"top_k_list": config.task.top_k_list}
 
     if config.task.type == "retrieval" and model_config is not None:
         task_kwargs.update(
@@ -68,19 +71,23 @@ def create_dataset(config: DatasetSourceConfig) -> Any:
         )
 
 
-def _get_table_name(config: BenchmarkConfig, model_name: str) -> str:
+def _get_table_name(
-    """Generate database table name from config and model name.
+    config: BenchmarkConfig,
+    model_name: str,
+    dataset_path: str,
+) -> str:
+    """Generate database table name from config, model, and dataset.
 
     Args:
         config: Benchmark configuration.
         model_name: Model name for table naming.
+        dataset_path: Dataset identifier or path (e.g. "uoft-cs/cifar100").
 
     Returns:
         Formatted table name.
     """
     prefix = config.model_table_prefix
-    # Use dataset path as part of table name (sanitize)
+    dataset_name = Path(dataset_path).name.lower().replace("-", "_")
-    dataset_name = Path(config.dataset.path).name.lower().replace("-", "_")
     return f"{prefix}_{dataset_name}_{model_name}"
 
 
@@ -88,6 +95,7 @@ def _ensure_table(
     config: BenchmarkConfig,
     model_name: str,
     vector_dim: int,
+    dataset_path: str,
 ) -> lancedb.table.Table:
     """Ensure the LanceDB table exists with correct schema.
 
@@ -95,6 +103,7 @@ def _ensure_table(
         config: Benchmark configuration.
         model_name: Model name for table naming.
         vector_dim: Feature vector dimension.
+        dataset_path: Dataset identifier or path.
 
     Returns:
         LanceDB table instance.
@@ -102,7 +111,7 @@ def _ensure_table(
     import pyarrow as pa
     from database import db_manager
 
-    table_name = _get_table_name(config, model_name)
+    table_name = _get_table_name(config, model_name, dataset_path)
 
     # Build expected schema
     schema = pa.schema(
@@ -132,7 +141,11 @@ def _ensure_table(
 
 
 def _print_benchmark_info(
-    config: BenchmarkConfig, vector_dim: int, table_name: str, table_count: int
+    config: BenchmarkConfig,
+    vector_dim: int,
+    table_name: str,
+    table_count: int,
+    dataset_path: str = "",
 ) -> None:
     """Print benchmark configuration info using Rich table.
 
@@ -141,12 +154,13 @@ def _print_benchmark_info(
         vector_dim: Feature vector dimension.
         table_name: Database table name.
         table_count: Number of entries in the table.
+        dataset_path: Current dataset identifier or path.
     """
     table = Table(title="Benchmark Configuration", show_header=False)
     table.add_column("Key", style="cyan", no_wrap=True)
     table.add_column("Value", style="magenta")
 
-    table.add_row("Dataset", f"{config.dataset.source_type} - {config.dataset.path}")
+    table.add_row("Dataset", dataset_path)
     table.add_row("Model Output Dimension", str(vector_dim))
     table.add_row("Table Name", table_name)
     table.add_row("Table Entries", str(table_count))
@@ -158,13 +172,23 @@ def _print_benchmark_results(results: dict[str, Any]) -> None:
     """Print benchmark results using Rich table.
 
     Args:
-        results: Final benchmark metrics.
+        results: Final benchmark metrics for a single dataset.
     """
     table = Table(title="Benchmark Results", show_header=False)
     table.add_column("Metric", style="cyan", no_wrap=True)
     table.add_column("Value", style="green")
 
+    # Print recalls first
+    recalls = results.get("recalls", {})
+    for k_str, v in recalls.items():
+        if isinstance(v, float):
+            table.add_row(k_str, f"{v:.4f}")
+
     for key, value in results.items():
+        if key == "recalls":
+            continue
+        if isinstance(value, (list, dict)):
+            continue
         if isinstance(value, float):
             table.add_row(key, f"{value:.4f}")
             continue
@@ -173,20 +197,147 @@ def _print_benchmark_results(results: dict[str, Any]) -> None:
     console.print(table)
 
 
+def _save_benchmark_outputs(
+    all_results: dict[str, dict[str, Any]],
+    output_root: Path,
+    model_name: str,
+) -> Path:
+    """Save multi-dataset benchmark results to disk.
+
+    Writes to ``output_root/{run_id}/``:
+        - summary.md — overall summary with per-dataset recall table
+        - metrics.csv — one row per (dataset, K) combination
+        - {dataset_name}/predictions.csv — per-sample predictions
+        - {dataset_name}/confusion_matrix.csv — row-normalized (Top-1)
+
+    Args:
+        all_results: Dict mapping dataset name to evaluate() result dict.
+        output_root: Parent directory (e.g. ``outputs/benchmark``).
+        model_name: Model identifier for the run_id.
+
+    Returns:
+        Path to the created run directory.
+    """
+    run_id = f"{datetime.now().strftime('%Y-%m-%d-%H%M%S')}-{model_name}"
+    out_dir = output_root / run_id
+    out_dir.mkdir(parents=True, exist_ok=True)
+
+    # --- metrics.csv: one row per (dataset, k) ---
+    csv_keys = ["model", "run_id", "dataset", "k", "recall@k", "total", "num_classes"]
+    all_csv_rows: list[dict[str, Any]] = []
+
+    # --- summary.md lines ---
+    md_lines = [
+        "# Benchmark Results",
+        "",
+        f"- **run_id**: `{run_id}`",
+        f"- **model**: `{model_name}`",
+        "",
+        "## Per-Dataset Recall",
+        "",
+    ]
+
+    for dataset_name, results in all_results.items():
+        recalls = results.get("recalls", {})
+        total = results.get("total", 0)
+        num_classes = results.get("num_classes", 0)
+
+        md_lines.append(f"### {dataset_name}")
+        md_lines.append("")
+        md_lines.append(f"- total queries: `{total}`")
+        md_lines.append(f"- num_classes: `{num_classes}`")
+        md_lines.append("")
+        md_lines.append("| K | Recall@K |")
+        md_lines.append("|---:|---:|")
+
+        for k_str, recall_val in sorted(recalls.items()):
+            k = int(k_str.replace("recall@", ""))
+            md_lines.append(f"| {k} | {recall_val:.4f} |")
+            all_csv_rows.append({
+                "model": model_name,
+                "run_id": run_id,
+                "dataset": dataset_name,
+                "k": k,
+                "recall@k": recall_val,
+                "total": total,
+                "num_classes": num_classes,
+            })
+
+        md_lines.append("")
+
+        # Per-dataset subdirectory
+        dataset_dir = out_dir / dataset_name
+        dataset_dir.mkdir(parents=True, exist_ok=True)
+
+        # --- predictions.csv (per-sample) ---
+        query_labels = results.get("query_labels", [])
+        topk_labels = results.get("topk_labels", [])
+        if query_labels and topk_labels:
+            max_k = max(len(preds) for preds in topk_labels) if topk_labels else 1
+            fieldnames = ["idx", "true_label"] + [
+                f"top{i+1}_label" for i in range(max_k)
+            ]
+            with (dataset_dir / "predictions.csv").open(
+                "w", newline="", encoding="utf-8"
+            ) as f:
+                writer = csv.DictWriter(f, fieldnames=fieldnames)
+                writer.writeheader()
+                for idx, (true_label, preds) in enumerate(
+                    zip(query_labels, topk_labels)
+                ):
+                    row_data: dict[str, Any] = {
+                        "idx": idx,
+                        "true_label": true_label,
+                    }
+                    for i in range(max_k):
+                        row_data[f"top{i+1}_label"] = (
+                            preds[i] if i < len(preds) else ""
+                        )
+                    writer.writerow(row_data)
+
+        # --- confusion_matrix.csv ---
+        cm = results.get("confusion_matrix")
+        if cm is not None:
+            with (dataset_dir / "confusion_matrix.csv").open(
+                "w", newline="", encoding="utf-8"
+            ) as fh:
+                writer = csv.writer(fh)
+                for row in cm:
+                    writer.writerow([f"{v:.6f}" for v in row])
+
+    # Write metrics.csv
+    if all_csv_rows:
+        with (out_dir / "metrics.csv").open("w", newline="", encoding="utf-8") as f:
+            writer = csv.DictWriter(f, fieldnames=csv_keys)
+            writer.writeheader()
+            for row in all_csv_rows:
+                writer.writerow(row)
+
+    # Write summary.md
+    (out_dir / "summary.md").write_text("\n".join(md_lines) + "\n", encoding="utf-8")
+
+    console.print(f"[green]Results saved to:[/green] {out_dir}")
+    return out_dir
+
+
 def run_benchmark(
     model: Any,
     processor: Any,
     config: BenchmarkConfig,
     model_config: ModelConfig | None = None,
     model_name: str = "model",
+    output_root: Path | None = None,
 ) -> dict[str, Any]:
-    """Run benchmark evaluation.
+    """Run benchmark evaluation across one or more datasets.
 
     Workflow:
-        1. Create dataset from configuration
+        1. Load model / processor once
-        2. Create benchmark task from configuration
+        2. For each dataset in config.datasets:
-        3. Build evaluation database from training set
+            a. Create dataset, get splits
-        4. Evaluate on test set
+            b. Build evaluation database
+            c. Evaluate with all K values in top_k_list
+        3. Print per-dataset results
+        4. Save all results to disk (if ``output_root`` provided)
 
     Args:
         model: Feature extraction model.
@@ -194,28 +345,15 @@ def run_benchmark(
         config: Benchmark configuration.
         model_config: Optional model configuration for task-owned loading.
         model_name: Model name for table naming.
+        output_root: Optional root directory for saving results.
 
     Returns:
-        Dictionary containing evaluation results.
+        Dict mapping dataset name to evaluation result dict.
 
     Raises:
         ValueError: If benchmark is not enabled in config.
     """
-    # Create dataset
+    # ── Resolve model / processor once ──────────────────────────
-    console.print(
-        f"[cyan]Loading dataset:[/cyan] {config.dataset.source_type} - {config.dataset.path}"
-    )
-    dataset = create_dataset(config.dataset)
-
-    # Get train and test splits
-    train_dataset = dataset.get_train_split()
-    test_dataset = dataset.get_test_split()
-
-    if train_dataset is None or test_dataset is None:
-        raise ValueError(
-            f"Dataset {config.dataset.path} does not have train/test splits"
-        )
-
     task = _create_task(config, model_config)
 
     resolver = getattr(task, "prepare_benchmark", None)
@@ -224,17 +362,16 @@ def run_benchmark(
             Callable[[Any, Any, str], tuple[Any, Any, str]],
             resolver,
         )
-        model, processor, model_name = prepare_benchmark(
+        model, processor, model_name = prepare_benchmark(model, processor, model_name)
-            model,
-            processor,
-            model_name,
-        )
 
     if model is None or processor is None:
         raise ValueError("Benchmark task did not provide a valid model and processor")
 
-    # Infer vector dimension from a sample
+    # ── Infer vector dimension from first dataset ─────────────────
-    sample = train_dataset[0]
+    first_cfg = config.datasets[0]
+    first_dataset = create_dataset(first_cfg)
+    first_train = first_dataset.get_train_split()
+    sample = first_train[0]
     sample_image = sample["img"]
 
     from utils.feature_extractor import infer_vector_dim
@@ -242,28 +379,77 @@ def run_benchmark(
     vector_dim = infer_vector_dim(processor, model, sample_image)
     console.print(f"[cyan]Model output dimension:[/cyan] {vector_dim}")
 
-    # Ensure table exists with correct schema
+    # ── Evaluate each dataset ────────────────────────────────────
-    table = _ensure_table(config, model_name, vector_dim)
+    all_results: dict[str, dict[str, Any]] = {}
-    table_name = _get_table_name(config, model_name)
 
-    # Check if database is already built
+    for dataset_cfg in config.datasets:
-    table_count = table.count_rows()
-    if table_count > 0:
         console.print(
-            f"[yellow]Table '{table_name}' already has {table_count} entries, skipping database build.[/yellow]"
+            f"\n[bold cyan]Dataset:[/bold cyan] "
+            f"{dataset_cfg.source_type} - {dataset_cfg.path}"
         )
-    else:
+
-        console.print(
+        dataset = create_dataset(dataset_cfg)
-            f"[cyan]Building database[/cyan] with {len(train_dataset)} training samples..."
+        train_dataset = dataset.get_train_split()
-        )
+        test_dataset = dataset.get_test_split()
-        task.build_database(model, processor, train_dataset, table, config.batch_size)
+
+        if train_dataset is None or test_dataset is None:
+            raise ValueError(
+                f"Dataset {dataset_cfg.path} does not have train/test splits"
+            )
+
+        dataset_name = Path(dataset_cfg.path).name.lower()
+        table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)
+        table_name = _get_table_name(config, model_name, dataset_cfg.path)
+
+        # Build database (skip if cached)
         table_count = table.count_rows()
+        expected_count = len(train_dataset)
 
-    _print_benchmark_info(config, vector_dim, table_name, table_count)
+        if table_count == expected_count:
+            console.print(
+                f"[yellow]Table '{table_name}' already has {table_count} entries, "
+                f"skipping database build.[/yellow]"
+            )
+        else:
+            if table_count > 0:
+                console.print(
+                    f"[yellow]Table '{table_name}' has {table_count} entries "
+                    f"(expected {expected_count}), rebuilding.[/yellow]"
+                )
+                # Rebuild: drop and recreate
+                from database import db_manager
+                db_manager.db.drop_table(table_name)
+                table = _ensure_table(config, model_name, vector_dim, dataset_cfg.path)
 
-    # Run evaluation (results with Rich table will be printed by the task)
+            console.print(
-    console.print(f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples...")
+                f"[cyan]Building database[/cyan] with {expected_count} training samples..."
-    results = task.evaluate(model, processor, test_dataset, table, config.batch_size)
+            )
-    _print_benchmark_results(results)
+            task.build_database(
+                model, processor, train_dataset, table, config.batch_size,
+                label_column=dataset_cfg.label_column,
+            )
+            table_count = table.count_rows()
 
-    return results
+        _print_benchmark_info(
+            config, vector_dim, table_name, table_count,
+            dataset_path=dataset_cfg.path,
+        )
+
+        # Evaluate
+        console.print(
+            f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples..."
+        )
+        results = task.evaluate(
+            model, processor, test_dataset, table, config.batch_size,
+            label_column=dataset_cfg.label_column,
+        )
+        all_results[dataset_name] = results
+
+        # Print per-dataset results
+        _print_benchmark_results(results)
+
+    # ── Save outputs ─────────────────────────────────────────────
+    if output_root is not None:
+        _save_benchmark_outputs(all_results, output_root, model_name)
+
+    return all_results

mini-nav/benchmarks/tasks/retrieval.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING, Any, cast
 
 import lancedb
+import numpy as np
 import pyarrow as pa
 import torch
 import torch.nn.functional as F
@@ -81,6 +82,7 @@ def _establish_eval_database(
     model: nn.Module,
     table: lancedb.table.Table,
     dataloader: DataLoader[Any],
+    label_column: str = "label",
 ) -> None:
     """Extract features from training images and store them in a database table.
 
@@ -89,6 +91,7 @@ def _establish_eval_database(
         model: Feature extraction model.
         table: LanceDB table to store features.
         dataloader: DataLoader for the training dataset.
+        label_column: Column name for labels in the batch dict.
     """
     # Extract all features using the utility function
     all_features = extract_batch_features(processor, model, dataloader)
@@ -97,7 +100,7 @@ def _establish_eval_database(
     # Store features to database
     global_idx = 0
     for batch in track(dataloader, description="Storing eval database"):
-        labels = batch[config.benchmark.dataset.label_column]
+        labels = batch[label_column]
         labels_list = labels.tolist()
         batch_size = len(labels_list)
 
@@ -120,7 +123,8 @@ def _evaluate_recall(
     table: lancedb.table.Table,
     dataloader: DataLoader[Any],
     top_k: int,
-) -> tuple[int, int]:
+    label_column: str = "label",
+) -> dict[str, Any]:
     """Evaluate Recall@K by searching the database for each test image.
 
     Args:
@@ -129,9 +133,15 @@ def _evaluate_recall(
         table: LanceDB table to search against.
         dataloader: DataLoader for the test dataset.
         top_k: Number of top results to retrieve.
+        label_column: Column name for labels in the batch dict.
 
     Returns:
-        A tuple of (correct_count, total_count).
+        A dict with keys:
+            - correct: Number of correct predictions
+            - total: Total number of test samples
+            - query_labels: True labels for each query (length N)
+            - topk_ids: IDs of top-K retrieved items (N x K)
+            - topk_labels: Labels of top-K retrieved items (N x K)
     """
     # Extract all features using the utility function
     all_features = extract_batch_features(processor, model, dataloader)
@@ -140,9 +150,12 @@ def _evaluate_recall(
     correct = 0
     total = 0
     feature_idx = 0
+    query_labels: list[int] = []
+    topk_ids: list[list[int]] = []
+    topk_labels: list[list[int]] = []
 
     for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
-        labels = batch[config.benchmark.dataset.label_column]
+        labels = batch[label_column]
         labels_list = labels.tolist()
 
         for j in range(len(labels_list)):
@@ -151,19 +164,67 @@ def _evaluate_recall(
 
             results = (
                 table.search(feature)
-                .select(["label", "_distance"])
+                .select(["id", "label", "_distance"])
                 .limit(top_k)
                 .to_polars()
             )
 
+            retrieved_ids = results["id"].to_list()
             retrieved_labels = results["label"].to_list()
             if true_label in retrieved_labels:
                 correct += 1
             total += 1
 
+            query_labels.append(true_label)
+            topk_ids.append(retrieved_ids)
+            topk_labels.append(retrieved_labels)
+
         feature_idx += len(labels_list)
 
-    return correct, total
+    return {
+        "correct": correct,
+        "total": total,
+        "query_labels": query_labels,
+        "topk_ids": topk_ids,
+        "topk_labels": topk_labels,
+    }
+
+
+def _build_confusion_matrix(
+    query_labels: list[int],
+    topk_labels: list[list[int]],
+    num_classes: int,
+) -> np.ndarray:
+    """Build row-normalized confusion matrix from Top-1 predictions.
+
+    For each query, the true label is compared against the first
+    retrieved label (Top-1). The resulting matrix is row-normalized
+    so that each row sums to 1.0, representing "given true class X,
+    what fraction of queries were predicted as each class Y".
+
+    Args:
+        query_labels: True labels for each query (length N).
+        topk_labels: Top-K retrieved labels for each query (N x K).
+        num_classes: Total number of label classes.
+
+    Returns:
+        Row-normalized confusion matrix as (num_classes, num_classes)
+        float64 array. Rows with zero queries remain all-zero.
+    """
+    cm = np.zeros((num_classes, num_classes), dtype=np.int64)
+    for y_true, top_labels in zip(query_labels, topk_labels):
+        y_pred = top_labels[0]  # Top-1
+        cm[int(y_true), int(y_pred)] += 1
+
+    # Row normalize: each row sums to 1.0
+    row_sums = cm.sum(axis=1, keepdims=True)
+    cm_norm = np.divide(
+        cm.astype(np.float64),
+        row_sums,
+        out=np.zeros_like(cm, dtype=np.float64),
+        where=row_sums > 0,
+    )
+    return cm_norm
 
 
 @RegisterTask("retrieval")
@@ -173,6 +234,7 @@ class RetrievalTask(BaseBenchmarkTask):
     def __init__(
         self,
         top_k: int = 10,
+        top_k_list: list[int] | None = None,
         dino_model: str = "facebook/dinov2-large",
         compression_dim: int = 512,
         compressor_path: str | None = None,
@@ -180,13 +242,22 @@ class RetrievalTask(BaseBenchmarkTask):
         """Initialize retrieval task.
 
         Args:
-            top_k: Number of top results to retrieve for recall calculation.
+            top_k: Maximum K for retrieval search.
+                Deprecated — prefer ``top_k_list``.
+            top_k_list: List of K values to evaluate (all derived from
+                a single max-K search). When not provided, ``[top_k]``
+                is used as fallback.
             dino_model: DINO model name used for feature extraction.
             compression_dim: Output dimension of the hash compressor.
             compressor_path: Optional path to trained hash compressor weights.
         """
-        super().__init__(top_k=top_k)
+        if top_k_list is not None:
-        self.top_k = top_k
+            top_k_list = sorted(set(top_k_list))
+        else:
+            top_k_list = [top_k]
+        super().__init__(top_k_list=top_k_list)
+        self.top_k_list: list[int] = top_k_list
+        self.max_k: int = max(self.top_k_list)
         self.dino_model = dino_model
         self.compression_dim = compression_dim
         self.compressor_path = compressor_path
@@ -251,6 +322,7 @@ class RetrievalTask(BaseBenchmarkTask):
         train_dataset: Any,
         table: lancedb.table.Table,
         batch_size: int,
+        label_column: str = "label",
     ) -> None:
         """Build the evaluation database from training data.
 
@@ -260,6 +332,7 @@ class RetrievalTask(BaseBenchmarkTask):
             train_dataset: Training dataset.
             table: LanceDB table to store features.
             batch_size: Batch size for DataLoader.
+            label_column: Column name for labels in the batch dict.
         """
         # Get a sample image to infer vector dimension
         sample = train_dataset[0]
@@ -282,7 +355,7 @@ class RetrievalTask(BaseBenchmarkTask):
             shuffle=False,
             num_workers=4,
         )
-        _establish_eval_database(processor, model, table, train_loader)
+        _establish_eval_database(processor, model, table, train_loader, label_column)
 
     def evaluate(
         self,
@@ -291,6 +364,7 @@ class RetrievalTask(BaseBenchmarkTask):
         test_dataset: Any,
         table: lancedb.table.Table,
         batch_size: int,
+        label_column: str = "label",
     ) -> dict[str, Any]:
         """Evaluate the model on the test dataset.
 
@@ -300,13 +374,17 @@ class RetrievalTask(BaseBenchmarkTask):
             test_dataset: Test dataset.
             table: LanceDB table to search against.
             batch_size: Batch size for DataLoader.
+            label_column: Column name for labels in the batch dict.
 
         Returns:
             Dictionary containing evaluation results with keys:
-                - accuracy: Recall@K accuracy (0.0 ~ 1.0)
+                - recalls: Dict of ``{"recall@K": value}`` for each K
-                - correct: Number of correct predictions
                 - total: Total number of test samples
-                - top_k: The K value used
+                - top_k_list: List of K values evaluated
+                - num_classes: Number of label classes
+                - confusion_matrix: Row-normalized confusion matrix (Top-1)
+                - query_labels: True labels for each query sample
+                - topk_labels: Top-K retrieved labels per query sample
         """
         test_loader = DataLoader(
             test_dataset.with_format("torch"),
@@ -314,15 +392,39 @@ class RetrievalTask(BaseBenchmarkTask):
             shuffle=False,
             num_workers=4,
         )
-        correct, total = _evaluate_recall(
+        eval_data = _evaluate_recall(
-            processor, model, table, test_loader, self.top_k
+            processor, model, table, test_loader, self.max_k, label_column
         )
 
-        accuracy = correct / total if total > 0 else 0.0
+        total = eval_data["total"]
+        query_labels = eval_data["query_labels"]
+        topk_labels = eval_data["topk_labels"]
+
+        # Compute Recall@k for each k in top_k_list
+        recalls: dict[str, float] = {}
+        for k in self.top_k_list:
+            hits = sum(
+                1
+                for true, preds in zip(query_labels, topk_labels)
+                if true in preds[:k]
+            )
+            recalls[f"recall@{k}"] = hits / total if total > 0 else 0.0
+
+        # Infer number of classes from the data
+        all_labels = query_labels + [
+            label for labels in topk_labels for label in labels
+        ]
+        num_classes = max(all_labels) + 1 if all_labels else 1
+
+        # Confusion matrix from Top-1 only
+        cm = _build_confusion_matrix(query_labels, topk_labels, num_classes)
 
         return {
-            "accuracy": accuracy,
+            "recalls": recalls,
-            "correct": correct,
             "total": total,
-            "top_k": self.top_k,
+            "top_k_list": self.top_k_list,
+            "num_classes": num_classes,
+            "confusion_matrix": cm.tolist(),
+            "query_labels": query_labels,
+            "topk_labels": topk_labels,
         }

mini-nav/commands/benchmark.py

@@ -1,3 +1,5 @@
+"""Benchmark CLI command."""
+
 import typer
 from commands import app
 
@@ -8,19 +10,59 @@ def benchmark(
     model_path: str | None = typer.Option(
         None, "--model", "-m", help="Path to compressor model weights"
     ),
+    dataset: list[str] | None = typer.Option(
+        None,
+        "--dataset",
+        "-d",
+        help=(
+            "Override datasets (HuggingFace path, repeatable). "
+            "Each value creates a single-entry dataset list."
+        ),
+    ),
+    top_k: list[int] | None = typer.Option(
+        None,
+        "--top-k",
+        "-k",
+        help="Override top-K values (repeatable). E.g. --top-k 1 --top-k 5 --top-k 10",
+    ),
 ):
+    """Run benchmark evaluation across configured datasets.
+
+    Without flags, runs all datasets and K values from config.yaml.
+    Use --dataset to override datasets, --top-k to override K values.
+    """
     from benchmarks import run_benchmark
     from configs import cfg_manager
+    from configs.models import DatasetSourceConfig
 
     config = cfg_manager.get()
     benchmark_cfg = config.benchmark
+
     if model_path:
         config.model.compressor_path = model_path
 
-    _ = run_benchmark(
+    # CLI overrides
+    if dataset:
+        benchmark_cfg.datasets = [
+            DatasetSourceConfig(
+                source_type="huggingface",
+                path=d,
+            )
+            for d in dataset
+        ]
+
+    if top_k:
+        benchmark_cfg.task.top_k_list = sorted(set(top_k))
+
+    model_name = "hash_compressor" if config.model.compressor_path else "dinov2"
+    output_root = config.output.directory / "benchmark"
+
+    results = run_benchmark(
         model=None,
         processor=None,
         config=benchmark_cfg,
         model_config=config.model,
-        model_name="hash_compressor" if config.model.compressor_path else "dinov2",
+        model_name=model_name,
+        output_root=output_root,
     )
+    return results

mini-nav/configs/config.yaml

@@ -22,14 +22,18 @@ dataset:
     seed: 42
 
 benchmark:
-    dataset:
+    datasets:
-        source_type: "huggingface"
+        - source_type: "huggingface"
-        path: "uoft-cs/cifar100"
+          path: "uoft-cs/cifar100"
-        img_column: "img"
+          img_column: "img"
-        label_column: "fine_label"
+          label_column: "fine_label"
+        - source_type: "huggingface"
+          path: "uoft-cs/cifar10"
+          img_column: "img"
+          label_column: "label"
     task:
         name: "recall_at_k"
         type: "retrieval"
-        top_k: 1
+        top_k_list: [1, 5, 10]
     batch_size: 64
     model_table_prefix: "benchmark"

mini-nav/configs/models.py

@@ -129,7 +129,24 @@ class BenchmarkTaskConfig(BaseModel):
 
     name: str = Field(default="recall_at_k", description="Task name")
     type: str = Field(default="retrieval", description="Task type")
-    top_k: int = Field(default=10, gt=0, description="Top K for recall evaluation")
+    top_k_list: list[int] = Field(
+        default=[1, 5, 10],
+        description="Top-K values to evaluate (all derived from a single max-K search)",
+    )
+
+    @property
+    def max_k(self) -> int:
+        """Maximum K for the underlying search; all values in top_k_list <= max_k."""
+        return max(self.top_k_list) if self.top_k_list else 1
+
+    @field_validator("top_k_list", mode="after")
+    @classmethod
+    def validate_top_k_list(cls, v: list[int]) -> list[int]:
+        if not v:
+            raise ValueError("top_k_list must contain at least one value")
+        if any(k <= 0 for k in v):
+            raise ValueError("top_k_list values must be positive")
+        return sorted(set(v))
 
     # Multi-object retrieval specific settings
     gamma: float = Field(
@@ -148,7 +165,10 @@ class BenchmarkConfig(BaseModel):
 
     model_config = ConfigDict(extra="ignore")
 
-    dataset: DatasetSourceConfig = Field(default_factory=DatasetSourceConfig)
+    datasets: list[DatasetSourceConfig] = Field(
+        default_factory=lambda: [DatasetSourceConfig()],
+        description="Dataset configurations to evaluate (supports multiple).",
+    )
     task: BenchmarkTaskConfig = Field(default_factory=BenchmarkTaskConfig)
     batch_size: int = Field(default=64, gt=0, description="Batch size for DataLoader")
     model_table_prefix: str = Field(