Skip to content

feat: Embedding Model Benchmarking #16

Description

@aGallea

Problem

The tool defaults to openai/clip-vit-base-patch32 for images and BAAI/bge-small-en-v1.5 for text. Users have no way to evaluate whether these are the best models for their specific data, or how different models would affect cluster quality. Model selection is currently guesswork.

Proposed Solution

Add a benchmark mode that indexes a sample of data with multiple models, computes clustering quality metrics, and presents a comparison table. This enables data-driven model selection.

Implementation Details

Backend

New file: embedding_cluster/benchmark.py

  • BenchmarkRunner class
    • Takes: CSV data, list of models to compare, sample size, num_clusters
    • For each model:
      1. Generate embeddings for the sample
      2. Run k-means clustering
      3. Compute quality metrics:
        • Silhouette Score — How well-separated clusters are [-1, 1]
        • Davies-Bouldin Index — Lower is better (cluster compactness vs separation)
        • Calinski-Harabasz Index — Higher is better (between-cluster vs within-cluster variance)
      4. Record timing (embedding generation time, clustering time)
      5. Record embedding dimensionality
    • Return: comparison table with all metrics per model
  • Handle model loading/unloading efficiently (don't hold all models in memory simultaneously)
  • Support both image and text models (detect from model name or explicit type parameter)

New file: embedding_cluster/server/routes/benchmark.py

  • POST /api/benchmark
    • Request:
      {
        "csv_path": "uploaded_file.csv",
        "models": [
          { "name": "openai/clip-vit-base-patch32", "type": "image", "field": "imageUrl" },
          { "name": "openai/clip-vit-large-patch14", "type": "image", "field": "imageUrl" },
          { "name": "BAAI/bge-small-en-v1.5", "type": "text", "field": "productDisplayName" },
          { "name": "all-MiniLM-L6-v2", "type": "text", "field": "productDisplayName" }
        ],
        "sample_size": 500,
        "num_clusters": 10
      }
    • Response: array of results per model with metrics, timing, and dimensionality
    • Long-running: use WebSocket for progress (model X of N, embedding Y of Z)
  • GET /api/benchmark/models — List predefined model presets with descriptions
    • Image: clip-vit-base-patch32, clip-vit-large-patch14, siglip-base-patch16-224
    • Text: bge-small-en-v1.5, bge-base-en-v1.5, all-MiniLM-L6-v2, all-mpnet-base-v2

Modify: embedding_cluster/server/models.py

  • Add BenchmarkRequest, ModelSpec, BenchmarkResult, BenchmarkResponse models

Modify: embedding_cluster/server/app.py

  • Register benchmark router

Frontend

New file: frontend/src/pages/BenchmarkPage.tsx

  • Step 1: Upload CSV or select existing uploaded file
  • Step 2: Select models to compare (checkboxes from predefined list + custom model name input)
  • Step 3: Configure sample size and cluster count
  • Step 4: Run benchmark with progress tracking
  • Step 5: Results table with sortable columns:
    | Model | Type | Dimensions | Silhouette | Davies-Bouldin | Calinski-Harabasz | Embed Time | Total Time |
    • Color-code best/worst values per metric
    • Highlight recommended model (best silhouette score)

New file: frontend/src/components/benchmark/ModelSelector.tsx

  • Predefined model checkboxes with descriptions
  • Custom model input field
  • Group by type (image vs text)

New file: frontend/src/components/benchmark/BenchmarkResults.tsx

  • Sortable results table
  • Color-coded cells (green=best, red=worst per column)
  • "Use this model" button that navigates to Index page with model pre-filled

Modify: frontend/src/api/benchmark.ts

  • New API functions: runBenchmark(), getModelPresets()

Modify: Navigation

  • Add "Benchmark" to the navigation bar

Predefined Model Presets

{
  "image_models": [
    { "name": "openai/clip-vit-base-patch32", "description": "CLIP base (fastest, 512d)" },
    { "name": "openai/clip-vit-large-patch14", "description": "CLIP large (better quality, 768d, slower)" },
    { "name": "google/siglip-base-patch16-224", "description": "SigLIP base (newer CLIP variant, 768d)" }
  ],
  "text_models": [
    { "name": "BAAI/bge-small-en-v1.5", "description": "BGE small (fast, 384d)" },
    { "name": "BAAI/bge-base-en-v1.5", "description": "BGE base (balanced, 768d)" },
    { "name": "sentence-transformers/all-MiniLM-L6-v2", "description": "MiniLM (lightweight, 384d)" },
    { "name": "sentence-transformers/all-mpnet-base-v2", "description": "MPNet (highest quality, 768d)" }
  ]
}

Testing Requirements

Full-coverage tests are required for all new code.

Backend Tests

New file: tests/test_benchmark.py

  • Test BenchmarkRunner with mocked models
    • Mock embedding generation to return deterministic vectors
    • Verify metrics are computed correctly
    • Verify timing is recorded
    • Test with synthetic data that has known cluster structure
  • Test benchmark endpoint:
    • Valid request returns results for all models
    • Invalid model name returns error
    • Sample size larger than dataset is handled
    • Progress WebSocket messages are sent
  • Test model presets endpoint returns predefined list
  • Test edge cases: single model, all same model, very small sample

Frontend Tests

  • BenchmarkPage: step-by-step flow works
  • ModelSelector: checkboxes toggle, custom input works
  • BenchmarkResults: table renders, sorting works, color coding applies
  • Navigation: benchmark link added

Acceptance Criteria

  • User can select multiple models to compare
  • Benchmark runs each model on a sample of the data
  • Results show silhouette score, Davies-Bouldin, Calinski-Harabasz, timing, and dimensions
  • Results table is sortable and color-coded
  • Progress tracking during benchmark execution
  • Predefined model presets are available with descriptions
  • User can navigate to Index page with selected model pre-filled
  • All new backend code has full test coverage
  • All new frontend components have tests

Notes

  • After implementation, update README.md to document the benchmarking feature (add to Features list, add Benchmark section to Usage)
  • Benchmarking downloads models — warn users about disk space and bandwidth
  • Consider caching benchmark results to avoid re-running
  • Sample-based benchmarking keeps computation feasible even for large datasets
  • All quality metrics are available from scikit-learn (existing dependency)

Metadata

Metadata

Assignees

No one assigned

    Labels

    backendBackend / Python changesenhancementNew feature or requestfrontendFrontend / React changespriority: 3-lowLower priority

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions