refactor(benchmark): delegate model loading to tasks and support CIFAR-100

- Extract model loading logic from benchmark CLI into task-owned prepare_benchmark
- Add RetrievalEncoder class wrapping DINO with optional hash compression
- Add accelerate dependency for device management  
- Switch dataset from CIFAR-10 to CIFAR-100 with fine_label column

mini-nav/benchmarks/runner.py

@@ -1,18 +1,42 @@
 """Benchmark runner for executing evaluations."""
 
 from pathlib import Path
-from typing import Any
+from typing import Any, Callable, cast
 
 import lancedb
 from benchmarks.datasets import HuggingFaceDataset, LocalDataset
 from benchmarks.tasks import get_task
-from configs.models import BenchmarkConfig, DatasetSourceConfig
+from configs.models import BenchmarkConfig, DatasetSourceConfig, ModelConfig
 from rich.console import Console
 from rich.table import Table
 
 console = Console()
 
 
+def _create_task(config: BenchmarkConfig, model_config: ModelConfig | None) -> Any:
+    """Create benchmark task with task-specific model settings.
+
+    Args:
+        config: Benchmark configuration.
+        model_config: Optional model configuration for task-owned loading.
+
+    Returns:
+        Benchmark task instance.
+    """
+    task_kwargs: dict[str, Any] = {"top_k": config.task.top_k}
+
+    if config.task.type == "retrieval" and model_config is not None:
+        task_kwargs.update(
+            {
+                "dino_model": model_config.dino_model,
+                "compression_dim": model_config.compression_dim,
+                "compressor_path": model_config.compressor_path,
+            }
+        )
+
+    return get_task(config.task.type, **task_kwargs)
+
+
 def create_dataset(config: DatasetSourceConfig) -> Any:
     """Create a dataset instance from configuration.
 
@@ -130,10 +154,30 @@ def _print_benchmark_info(
     console.print(table)
 
 
+def _print_benchmark_results(results: dict[str, Any]) -> None:
+    """Print benchmark results using Rich table.
+
+    Args:
+        results: Final benchmark metrics.
+    """
+    table = Table(title="Benchmark Results", show_header=False)
+    table.add_column("Metric", style="cyan", no_wrap=True)
+    table.add_column("Value", style="green")
+
+    for key, value in results.items():
+        if isinstance(value, float):
+            table.add_row(key, f"{value:.4f}")
+            continue
+        table.add_row(key, str(value))
+
+    console.print(table)
+
+
 def run_benchmark(
     model: Any,
     processor: Any,
     config: BenchmarkConfig,
+    model_config: ModelConfig | None = None,
     model_name: str = "model",
 ) -> dict[str, Any]:
     """Run benchmark evaluation.
@@ -148,6 +192,7 @@ def run_benchmark(
         model: Feature extraction model.
         processor: Image preprocessor.
         config: Benchmark configuration.
+        model_config: Optional model configuration for task-owned loading.
         model_name: Model name for table naming.
 
     Returns:
@@ -171,6 +216,23 @@ def run_benchmark(
             f"Dataset {config.dataset.path} does not have train/test splits"
         )
 
+    task = _create_task(config, model_config)
+
+    resolver = getattr(task, "prepare_benchmark", None)
+    if callable(resolver):
+        prepare_benchmark = cast(
+            Callable[[Any, Any, str], tuple[Any, Any, str]],
+            resolver,
+        )
+        model, processor, model_name = prepare_benchmark(
+            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
     sample = train_dataset[0]
     sample_image = sample["img"]
@@ -191,16 +253,17 @@ def run_benchmark(
             f"[yellow]Table '{table_name}' already has {table_count} entries, skipping database build.[/yellow]"
         )
     else:
-        # Create and run benchmark task
-        task = get_task(config.task.type, top_k=config.task.top_k)
         console.print(
             f"[cyan]Building database[/cyan] with {len(train_dataset)} training samples..."
         )
         task.build_database(model, processor, train_dataset, table, config.batch_size)
+        table_count = table.count_rows()
+
+    _print_benchmark_info(config, vector_dim, table_name, table_count)
 
     # Run evaluation (results with Rich table will be printed by the task)
-    task = get_task(config.task.type, top_k=config.task.top_k)
     console.print(f"[cyan]Evaluating[/cyan] on {len(test_dataset)} test samples...")
     results = task.evaluate(model, processor, test_dataset, table, config.batch_size)
+    _print_benchmark_results(results)
 
     return results

mini-nav/benchmarks/tasks/retrieval.py

@@ -1,17 +1,62 @@
 """Retrieval task for benchmark evaluation (Recall@K)."""
 
-from typing import Any
+from typing import TYPE_CHECKING, Any, cast
 
 import lancedb
 import pyarrow as pa
+import torch
+import torch.nn.functional as F
 from benchmarks.base import BaseBenchmarkTask
 from benchmarks.tasks.registry import RegisterTask
+from compressors.model_loader import get_dino_dim, load_dino_model, load_hash_compressor
+from configs import cfg_manager
 from rich.progress import track
 from torch import nn
 from torch.utils.data import DataLoader
 from transformers import BitImageProcessor
 from utils.feature_extractor import extract_batch_features, infer_vector_dim
 
+if TYPE_CHECKING:
+    from compressors.hash_compressor import HashCompressor
+
+
+class RetrievalEncoder(nn.Module):
+    """Benchmark encoder for DINO and optional hash compression."""
+
+    def __init__(
+        self,
+        dino: nn.Module,
+        compressor: "HashCompressor | None" = None,
+    ) -> None:
+        """Initialize retrieval encoder.
+
+        Args:
+            dino: DINO backbone used for feature extraction.
+            compressor: Optional hash compressor for recall evaluation.
+        """
+        super().__init__()
+        self.dino: nn.Module = dino
+        self.compressor: HashCompressor | None = compressor
+
+    def forward(self, inputs: Any) -> torch.Tensor:
+        """Encode processor inputs into benchmark vectors.
+
+        Args:
+            inputs: Batched processor outputs.
+
+        Returns:
+            Float tensor used for LanceDB insertion and retrieval.
+        """
+        outputs = self.dino(**inputs)
+        tokens = outputs.last_hidden_state
+
+        if self.compressor is None:
+            features = tokens.mean(dim=1)
+            return F.normalize(features, dim=-1)
+
+        bits = self.compressor.encode(tokens)
+        return bits.to(dtype=torch.float32)
+
 
 def _build_eval_schema(vector_dim: int) -> pa.Schema:
     """Build PyArrow schema for evaluation database table.
@@ -35,7 +80,7 @@ def _establish_eval_database(
     processor: BitImageProcessor,
     model: nn.Module,
     table: lancedb.table.Table,
-    dataloader: DataLoader,
+    dataloader: DataLoader[Any],
 ) -> None:
     """Extract features from training images and store them in a database table.
 
@@ -47,11 +92,12 @@ def _establish_eval_database(
     """
     # Extract all features using the utility function
     all_features = extract_batch_features(processor, model, dataloader)
+    config = cfg_manager.get()
 
     # Store features to database
     global_idx = 0
     for batch in track(dataloader, description="Storing eval database"):
-        labels = batch["label"]
+        labels = batch[config.benchmark.dataset.label_column]
         labels_list = labels.tolist()
         batch_size = len(labels_list)
 
@@ -72,7 +118,7 @@ def _evaluate_recall(
     processor: BitImageProcessor,
     model: nn.Module,
     table: lancedb.table.Table,
-    dataloader: DataLoader,
+    dataloader: DataLoader[Any],
     top_k: int,
 ) -> tuple[int, int]:
     """Evaluate Recall@K by searching the database for each test image.
@@ -89,13 +135,14 @@ def _evaluate_recall(
     """
     # Extract all features using the utility function
     all_features = extract_batch_features(processor, model, dataloader)
+    config = cfg_manager.get()
 
     correct = 0
     total = 0
     feature_idx = 0
 
     for batch in track(dataloader, description=f"Evaluating Recall@{top_k}"):
-        labels = batch["label"]
+        labels = batch[config.benchmark.dataset.label_column]
         labels_list = labels.tolist()
 
         for j in range(len(labels_list)):
@@ -123,14 +170,79 @@ def _evaluate_recall(
 class RetrievalTask(BaseBenchmarkTask):
     """Retrieval evaluation task (Recall@K)."""
 
-    def __init__(self, top_k: int = 10):
+    def __init__(
+        self,
+        top_k: int = 10,
+        dino_model: str = "facebook/dinov2-large",
+        compression_dim: int = 512,
+        compressor_path: str | None = None,
+    ):
         """Initialize retrieval task.
 
         Args:
             top_k: Number of top results to retrieve for recall calculation.
+            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)
         self.top_k = top_k
+        self.dino_model = dino_model
+        self.compression_dim = compression_dim
+        self.compressor_path = compressor_path
+        self._processor: BitImageProcessor | None = None
+        self._model: nn.Module | None = None
+        self._model_name = "hash_compressor" if compressor_path else "dinov2"
+
+    def prepare_benchmark(
+        self,
+        model: Any,
+        processor: Any,
+        model_name: str = "model",
+    ) -> tuple[nn.Module, BitImageProcessor, str]:
+        """Resolve benchmark resources for this task.
+
+        Args:
+            model: Optional pre-built model from the caller.
+            processor: Optional pre-built processor from the caller.
+            model_name: Fallback table model name.
+
+        Returns:
+            Tuple of benchmark model, processor, and resolved model name.
+        """
+        if model is not None and processor is not None:
+            return (
+                cast(nn.Module, model),
+                cast(BitImageProcessor, processor),
+                model_name,
+            )
+
+        self._ensure_resources_loaded()
+        return (
+            cast(nn.Module, self._model),
+            cast(BitImageProcessor, self._processor),
+            self._model_name,
+        )
+
+    def _ensure_resources_loaded(self) -> None:
+        """Lazy-load retrieval benchmark resources."""
+        if self._processor is not None and self._model is not None:
+            return
+
+        processor, dino = load_dino_model(self.dino_model)
+
+        compressor = None
+        if self.compressor_path is not None:
+            compressor = load_hash_compressor(
+                input_dim=get_dino_dim(self.dino_model),
+                hash_bits=self.compression_dim,
+                compressor_path=self.compressor_path,
+            )
+            compressor.eval()
+
+        self._processor = processor
+        self._model = RetrievalEncoder(dino=dino, compressor=compressor)
+        self._model.eval()
 
     def build_database(
         self,

mini-nav/commands/benchmark.py

@@ -1,79 +1,26 @@
-from typing import Any, Optional, cast
-
 import typer
 from commands import app
 
 
 @app.command()
 def benchmark(
-    ctx: typer.Context,
+    _ctx: typer.Context,
-    model_path: Optional[str] = typer.Option(
+    model_path: str | None = typer.Option(
         None, "--model", "-m", help="Path to compressor model weights"
     ),
 ):
-    import torch
-    import torch.nn.functional as F
     from benchmarks import run_benchmark
     from configs import cfg_manager
-    from transformers import AutoImageProcessor, AutoModel, BitImageProcessor
-    from utils import get_device
 
     config = cfg_manager.get()
     benchmark_cfg = config.benchmark
-
-    device = get_device()
-
-    model_cfg = config.model
-    processor = cast(
-        BitImageProcessor,
-        AutoImageProcessor.from_pretrained(model_cfg.dino_model, device_map=device),
-    )
-
-    # Load DINO model for feature extraction
-    dino = AutoModel.from_pretrained(model_cfg.dino_model, device_map=device)
-    dino.eval()
-
-    # Optional hash compressor
-    compressor = None
     if model_path:
-        from compressors import HashCompressor
+        config.model.compressor_path = model_path
 
-        compressor = HashCompressor(
+    _ = run_benchmark(
-            input_dim=model_cfg.compression_dim,
+        model=None,
-            hash_bits=model_cfg.compression_dim,
+        processor=None,
-        )
-        compressor.load_state_dict(torch.load(model_path))
-        compressor.to(device)
-        compressor.eval()
-
-    # Create wrapper with extract_features method
-    class DinoFeatureExtractor:
-        def __init__(self, dino, compressor=None):
-            self.dino = dino
-            self.compressor = compressor
-
-        def extract_features(self, images: list) -> torch.Tensor:
-            inputs = processor(images, return_tensors="pt").to(device)
-            with torch.no_grad():
-                outputs = self.dino(**inputs)
-                features = outputs.last_hidden_state.mean(dim=1)
-                features = F.normalize(features, dim=-1)
-            return features
-
-        def encode(self, images: list) -> torch.Tensor:
-            if self.compressor is None:
-                return self.extract_features(images)
-            tokens = self.dino(
-                **processor(images, return_tensors="pt").to(device)
-            ).last_hidden_state
-            _, _, bits = self.compressor(tokens)
-            return bits
-
-    model = DinoFeatureExtractor(dino, compressor)
-
-    run_benchmark(
-        model=model,
-        processor=processor,
         config=benchmark_cfg,
-        model_name="dinov2",
+        model_config=config.model,
+        model_name="hash_compressor" if config.model.compressor_path else "dinov2",
     )

mini-nav/configs/config.yaml

@@ -24,9 +24,9 @@ dataset:
 benchmark:
     dataset:
         source_type: "huggingface"
-        path: "uoft-cs/cifar10"
+        path: "uoft-cs/cifar100"
         img_column: "img"
-        label_column: "label"
+        label_column: "fine_label"
     task:
         name: "recall_at_k"
         type: "retrieval"

mini-nav/utils/feature_extractor.py

@@ -3,11 +3,13 @@
 from typing import Any, List, Union, cast
 
 import torch
+from aiohttp.web import get
 from PIL import Image
+from rich.progress import track
 from torch import nn
 from torch.utils.data import DataLoader
 from transformers import BitImageProcessor
-from rich.progress import track
+from utils import get_device
 
 
 def _extract_features_from_output(output: Any) -> torch.Tensor:
@@ -26,6 +28,7 @@ def _extract_features_from_output(output: Any) -> torch.Tensor:
     return cast(torch.Tensor, output)
 
 
+@torch.no_grad()
 def infer_vector_dim(
     processor: BitImageProcessor,
     model: nn.Module,
@@ -41,12 +44,8 @@ def infer_vector_dim(
     Returns:
         Vector dimension.
     """
-    device = next(model.parameters()).device
-    model.eval()
-
-    with torch.no_grad():
     inputs = processor(images=sample_image, return_tensors="pt")
-        inputs.to(device)
+    inputs.to(get_device())
     output = model(inputs)
 
     features = _extract_features_from_output(output)
@@ -69,11 +68,8 @@ def extract_single_image_feature(
     Returns:
         The extracted CLS token feature vector as a list of floats.
     """
-    device = next(model.parameters()).device
-    model.eval()
-
     inputs = processor(images=image, return_tensors="pt")
-    inputs.to(device, non_blocking=True)
+    inputs.to(get_device(), non_blocking=True)
     outputs = model(inputs)
 
     features = _extract_features_from_output(outputs)  # [1, D]
@@ -98,8 +94,7 @@ def extract_batch_features(
     Returns:
         Tensor of shape [batch_size, feature_dim].
     """
-    device = next(model.parameters()).device
+    device = get_device()
-    model.eval()
 
     # Handle DataLoader input
     if isinstance(images, DataLoader):

pyproject.toml

@@ -5,6 +5,7 @@ description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.10"
 dependencies = [
+    "accelerate>=1.13.0",
     "altair>=6.0.0",
     "dash>=3.4.0",
     "dash-ag-grid>=33.3.3",