init

README.md

@@ -1,2 +1,69 @@
 # embedding-bench
-Compare text embedding models across retrieval performance, inference speed, and memory footprint.
+
+Compare text embedding models across retrieval quality, inference speed, and memory footprint. Everything runs locally — no external API calls.
+
+## Models
+
+| Key | Model | Role |
+|-----|-------|------|
+| `mpnet` | `sentence-transformers/all-mpnet-base-v2` | Baseline |
+| `bge-small` | `BAAI/bge-small-en-v1.5` | |
+
+## Setup
+
+```bash
+python3 -m venv .venv
+source .venv/bin/activate
+pip install -r requirements.txt
+```
+
+## Usage
+
+```bash
+# Full benchmark (quality + speed + memory)
+python bench.py
+
+# Specific models
+python bench.py --models mpnet bge-small
+
+# Skip expensive evals
+python bench.py --skip-quality
+python bench.py --skip-memory
+
+# Tune corpus size and batch size
+python bench.py --corpus-size 500 --batch-size 32 --num-runs 5
+```
+
+## Metrics
+
+| Dimension | Metric | Method |
+|-----------|--------|--------|
+| Quality | Spearman rho | STS Benchmark test set (1,379 pairs) |
+| Speed | Median encode time | Wall-clock over N runs with warmup |
+| Memory | Peak RSS delta | Isolated subprocess via `psutil` |
+
+## Adding a model
+
+Edit `models.py` and add an entry to `REGISTRY`:
+
+```python
+"e5-small": ModelConfig(
+    name="e5-small-v2",
+    model_id="intfloat/e5-small-v2",
+),
+```
+
+## Project structure
+
+```
+embedding-bench/
+├── bench.py           # CLI entry point
+├── models.py          # Model registry
+├── corpus.py          # Sentence corpus builder
+├── report.py          # Table formatting
+├── evals/
+│   ├── quality.py     # STS Benchmark evaluation
+│   ├── speed.py       # Latency measurement
+│   └── memory.py      # Memory measurement
+└── requirements.txt
+```

bench.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+
+import argparse
+
+from sentence_transformers import SentenceTransformer
+
+from corpus import build_corpus
+from evals import evaluate_memory, evaluate_quality, evaluate_speed
+from models import REGISTRY
+from report import print_report
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(
+        prog="embedding-bench",
+        description="Compare embedding models on quality, speed, and memory.",
+    )
+    parser.add_argument(
+        "--models",
+        nargs="+",
+        default=list(REGISTRY.keys()),
+        choices=list(REGISTRY.keys()),
+        help="Models to benchmark (default: all)",
+    )
+    parser.add_argument("--corpus-size", type=int, default=1000)
+    parser.add_argument("--batch-size", type=int, default=64)
+    parser.add_argument("--num-runs", type=int, default=3)
+    parser.add_argument("--skip-quality", action="store_true")
+    parser.add_argument("--skip-speed", action="store_true")
+    parser.add_argument("--skip-memory", action="store_true")
+
+    args = parser.parse_args(argv)
+
+    configs = [REGISTRY[k] for k in args.models]
+    baseline_name = next((c.name for c in configs if c.is_baseline), None)
+
+    corpus: list[str] | None = None
+    if not args.skip_speed or not args.skip_memory:
+        print(f"Preparing corpus ({args.corpus_size} sentences)...")
+        corpus = build_corpus(args.corpus_size)
+
+    results = []
+    for cfg in configs:
+        print(f"\n{'='*50}")
+        print(f"Benchmarking: {cfg.name}")
+        print(f"{'='*50}")
+
+        result: dict = {"name": cfg.name, "is_baseline": cfg.is_baseline}
+
+        if not args.skip_quality:
+            print("  Evaluating quality (STS Benchmark)...")
+            model = SentenceTransformer(cfg.model_id)
+            result["quality"] = evaluate_quality(model)
+            print(f"  Quality: {result['quality']:.4f}")
+            del model
+
+        if not args.skip_speed and corpus is not None:
+            print(f"  Evaluating speed ({args.num_runs} runs, {args.corpus_size} sentences)...")
+            model = SentenceTransformer(cfg.model_id)
+            result["speed"] = evaluate_speed(model, corpus, num_runs=args.num_runs, batch_size=args.batch_size)
+            print(f"  Speed: {result['speed']['sentences_per_second']} sent/s")
+            del model
+
+        if not args.skip_memory and corpus is not None:
+            print("  Evaluating memory (isolated subprocess)...")
+            result["memory_mb"] = evaluate_memory(cfg.model_id, corpus, batch_size=args.batch_size)
+            print(f"  Memory: {result['memory_mb']} MB")
+
+        results.append(result)
+
+    print_report(results, baseline_name=baseline_name)
+
+
+if __name__ == "__main__":
+    main()

corpus.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from datasets import load_dataset
+
+
+def build_corpus(size: int) -> list[str]:
+    """Build a corpus of real sentences from the STS Benchmark dataset."""
+    dataset = load_dataset("mteb/stsbenchmark-sts", split="test")
+    sentences = list(dataset["sentence1"]) + list(dataset["sentence2"])
+    full: list[str] = []
+    while len(full) < size:
+        full.extend(sentences)
+    return full[:size]

evals/__init__.py

@@ -0,0 +1,5 @@
+from evals.quality import evaluate_quality
+from evals.speed import evaluate_speed
+from evals.memory import evaluate_memory
+
+__all__ = ["evaluate_quality", "evaluate_speed", "evaluate_memory"]

evals/memory.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+import multiprocessing
+import os
+
+
+def _measure(model_id: str, sentences: list[str], batch_size: int, queue: multiprocessing.Queue) -> None:
+    import psutil
+    from sentence_transformers import SentenceTransformer
+
+    process = psutil.Process(os.getpid())
+    baseline = process.memory_info().rss
+    model = SentenceTransformer(model_id)
+    model.encode(sentences, batch_size=batch_size, show_progress_bar=False)
+    peak = process.memory_info().rss
+    queue.put(peak - baseline)
+
+
+def evaluate_memory(model_id: str, sentences: list[str], batch_size: int = 64) -> float:
+    """Return memory delta in MB, measured in an isolated subprocess."""
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    p = ctx.Process(target=_measure, args=(model_id, sentences, batch_size, q))
+    p.start()
+    p.join()
+    bytes_delta = q.get()
+    return round(bytes_delta / (1024 * 1024), 1)

evals/quality.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from datasets import load_dataset
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
+
+
+def evaluate_quality(model: SentenceTransformer) -> float:
+    """Return Spearman correlation on the STS Benchmark test set."""
+    dataset = load_dataset("mteb/stsbenchmark-sts", split="test")
+    sentences1 = list(dataset["sentence1"])
+    sentences2 = list(dataset["sentence2"])
+    scores = [s / 5.0 for s in dataset["score"]]
+
+    evaluator = EmbeddingSimilarityEvaluator(sentences1, sentences2, scores)
+    results = evaluator(model)
+    return results["spearman_cosine"]

evals/speed.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import statistics
+import time
+
+from sentence_transformers import SentenceTransformer
+
+
+def evaluate_speed(
+    model: SentenceTransformer,
+    sentences: list[str],
+    num_runs: int = 3,
+    batch_size: int = 64,
+) -> dict[str, float]:
+    """Measure encoding latency. Returns median time and throughput."""
+    model.encode(sentences, batch_size=batch_size, show_progress_bar=False)
+
+    times: list[float] = []
+    for _ in range(num_runs):
+        start = time.perf_counter()
+        model.encode(sentences, batch_size=batch_size, show_progress_bar=False)
+        elapsed = time.perf_counter() - start
+        times.append(elapsed)
+
+    median_time = statistics.median(times)
+    return {
+        "median_seconds": round(median_time, 4),
+        "sentences_per_second": round(len(sentences) / median_time, 1),
+    }

models.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass
+class ModelConfig:
+    name: str
+    model_id: str
+    is_baseline: bool = False
+
+
+REGISTRY: dict[str, ModelConfig] = {
+    "mpnet": ModelConfig(
+        name="all-mpnet-base-v2",
+        model_id="sentence-transformers/all-mpnet-base-v2",
+        is_baseline=True,
+    ),
+    "bge-small": ModelConfig(
+        name="bge-small-en-v1.5",
+        model_id="BAAI/bge-small-en-v1.5",
+    ),
+}

report.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from tabulate import tabulate
+
+
+def print_report(results: list[dict[str, Any]], baseline_name: Optional[str] = None) -> None:
+    """Print a formatted comparison table to stdout."""
+    headers = ["Model", "Quality (STS)", "Speed (sent/s)", "Median Time (s)", "Memory (MB)"]
+    rows: list[list[Any]] = []
+
+    for r in results:
+        name = r["name"]
+        if r.get("is_baseline"):
+            name += " [B]"
+
+        quality = r.get("quality")
+        speed = r.get("speed")
+        memory = r.get("memory_mb")
+
+        rows.append([
+            name,
+            f"{quality:.4f}" if quality is not None else "—",
+            f"{speed['sentences_per_second']}" if speed else "—",
+            f"{speed['median_seconds']}" if speed else "—",
+            f"{memory}" if memory is not None else "—",
+        ])
+
+    print()
+    print(tabulate(rows, headers=headers, tablefmt="simple"))
+    if baseline_name:
+        print(f"\n[B] = baseline ({baseline_name})")
+    print()

requirements.txt

@@ -0,0 +1,5 @@
+sentence-transformers>=2.2.0
+torch
+datasets
+psutil
+tabulate