add fastembed lib

README.md

@@ -4,10 +4,18 @@ Compare text embedding models across retrieval quality, inference speed, and mem
 
 ## Models
 
-| Key | Model | Role |
-|-----|-------|------|
-| `mpnet` | `sentence-transformers/all-mpnet-base-v2` | Baseline |
-| `bge-small` | `BAAI/bge-small-en-v1.5` | |
+| Key | Model | Backend | Role |
+|-----|-------|---------|------|
+| `mpnet` | `sentence-transformers/all-mpnet-base-v2` | sbert | Baseline |
+| `bge-small` | `BAAI/bge-small-en-v1.5` | sbert | |
+| `bge-small-fe` | `BAAI/bge-small-en-v1.5` | fastembed | |
+| `all-minilm-fe` | `sentence-transformers/all-MiniLM-L6-v2` | fastembed | |
+
+Three backends are supported:
+
+- **sbert** — [sentence-transformers](https://www.sbert.net/) (PyTorch). Default.
+- **fastembed** — [qdrant/fastembed](https://github.com/qdrant/fastembed) (ONNX Runtime). Lighter and often faster.
+- **gguf** — [llama-cpp-python](https://github.com/abetlen/llama-cpp-python) for quantised GGUF models.
 
 ## Setup
 
@@ -26,6 +34,9 @@ python bench.py
 # Specific models
 python bench.py --models mpnet bge-small
 
+# Compare the same model across backends
+python bench.py --models bge-small bge-small-fe
+
 # Skip expensive evals
 python bench.py --skip-quality
 python bench.py --skip-memory
@@ -47,10 +58,18 @@ python bench.py --corpus-size 500 --batch-size 32 --num-runs 5
 Edit `models.py` and add an entry to `REGISTRY`:
 
 ```python
+# sentence-transformers backend (default)
 "e5-small": ModelConfig(
     name="e5-small-v2",
     model_id="intfloat/e5-small-v2",
 ),
+
+# fastembed backend
+"e5-small-fe": ModelConfig(
+    name="e5-small-v2 (fastembed)",
+    model_id="intfloat/e5-small-v2",
+    backend="fastembed",
+),
 ```
 
 ## Project structure
@@ -59,6 +78,7 @@ Edit `models.py` and add an entry to `REGISTRY`:
 embedding-bench/
 ├── bench.py           # CLI entry point
 ├── models.py          # Model registry
+├── wrapper.py         # Backend wrappers (sbert, fastembed, gguf)
 ├── corpus.py          # Sentence corpus builder
 ├── report.py          # Table formatting
 ├── evals/

bench.py

@@ -2,12 +2,11 @@ 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
+from wrapper import load_model
 
 
 def main(argv: list[str] | None = None) -> None:
@@ -49,21 +48,21 @@ def main(argv: list[str] | None = None) -> None:
 
         if not args.skip_quality:
             print("  Evaluating quality (STS Benchmark)...")
-            model = SentenceTransformer(cfg.model_id)
+            model = load_model(cfg)
             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)
+            model = load_model(cfg)
             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)
+            result["memory_mb"] = evaluate_memory(cfg.model_id, corpus, batch_size=args.batch_size, backend=cfg.backend)
             print(f"  Memory: {result['memory_mb']} MB")
 
         results.append(result)

evals/memory.py

@@ -4,23 +4,25 @@ import multiprocessing
 import os
 
 
-def _measure(model_id: str, sentences: list[str], batch_size: int, queue: multiprocessing.Queue) -> None:
+def _measure(model_id: str, backend: str, sentences: list[str], batch_size: int, queue: multiprocessing.Queue) -> None:
     import psutil
-    from sentence_transformers import SentenceTransformer
+    from models import ModelConfig
+    from wrapper import load_model
 
     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)
+    cfg = ModelConfig(name="", model_id=model_id, backend=backend)
+    model = load_model(cfg)
+    model.encode(sentences, batch_size=batch_size)
     peak = process.memory_info().rss
     queue.put(peak - baseline)
 
 
-def evaluate_memory(model_id: str, sentences: list[str], batch_size: int = 64) -> float:
+def evaluate_memory(model_id: str, sentences: list[str], batch_size: int = 64, backend: str = "sbert") -> 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 = ctx.Process(target=_measure, args=(model_id, backend, sentences, batch_size, q))
     p.start()
     p.join()
     bytes_delta = q.get()

evals/quality.py

@@ -1,17 +1,24 @@
 from __future__ import annotations
 
+import numpy as np
 from datasets import load_dataset
-from sentence_transformers import SentenceTransformer
-from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
+from scipy.stats import spearmanr
 
 
-def evaluate_quality(model: SentenceTransformer) -> float:
+def evaluate_quality(model) -> 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"]]
+    gold_scores = [s / 5.0 for s in dataset["score"]]
 
-    evaluator = EmbeddingSimilarityEvaluator(sentences1, sentences2, scores)
-    results = evaluator(model)
-    return results["spearman_cosine"]
+    emb1 = model.encode(sentences1)
+    emb2 = model.encode(sentences2)
+
+    # Row-wise cosine similarity
+    cos_sims = np.sum(emb1 * emb2, axis=1) / (
+        np.linalg.norm(emb1, axis=1) * np.linalg.norm(emb2, axis=1)
+    )
+
+    correlation, _ = spearmanr(cos_sims, gold_scores)
+    return correlation

evals/speed.py

@@ -3,11 +3,9 @@ from __future__ import annotations
 import statistics
 import time
 
-from sentence_transformers import SentenceTransformer
-
 
 def evaluate_speed(
-    model: SentenceTransformer,
+    model,
     sentences: list[str],
     num_runs: int = 3,
     batch_size: int = 64,

models.py

@@ -8,6 +8,8 @@ class ModelConfig:
     name: str
     model_id: str
     is_baseline: bool = False
+    backend: str = "sbert"
+    gguf_file: str | None = None
 
 
 REGISTRY: dict[str, ModelConfig] = {
@@ -20,4 +22,14 @@ REGISTRY: dict[str, ModelConfig] = {
         name="bge-small-en-v1.5",
         model_id="BAAI/bge-small-en-v1.5",
     ),
+    "bge-small-fe": ModelConfig(
+        name="bge-small-en-v1.5 (fastembed)",
+        model_id="BAAI/bge-small-en-v1.5",
+        backend="fastembed",
+    ),
+    "all-minilm-fe": ModelConfig(
+        name="all-MiniLM-L6-v2 (fastembed)",
+        model_id="sentence-transformers/all-MiniLM-L6-v2",
+        backend="fastembed",
+    ),
 }

requirements.txt

@@ -3,3 +3,6 @@ torch
 datasets
 psutil
 tabulate
+fastembed
+numpy
+scipy

wrapper.py

@@ -0,0 +1,60 @@
+from __future__ import annotations
+
+import numpy as np
+
+from models import ModelConfig
+
+
+class SBertWrapper:
+    """Wraps sentence_transformers.SentenceTransformer."""
+
+    def __init__(self, cfg: ModelConfig):
+        from sentence_transformers import SentenceTransformer
+        self._model = SentenceTransformer(cfg.model_id)
+
+    def encode(self, sentences: list[str], batch_size: int = 64, **kwargs) -> np.ndarray:
+        kwargs.setdefault("show_progress_bar", False)
+        return self._model.encode(sentences, batch_size=batch_size, **kwargs)
+
+
+class GGUFWrapper:
+    """Wraps llama_cpp.Llama in embedding mode."""
+
+    def __init__(self, cfg: ModelConfig):
+        from huggingface_hub import hf_hub_download
+        from llama_cpp import Llama
+
+        path = hf_hub_download(repo_id=cfg.model_id, filename=cfg.gguf_file)
+        self._model = Llama(
+            model_path=path, embedding=True, n_ctx=512, verbose=False
+        )
+
+    def encode(self, sentences: list[str], batch_size: int = 64, **kwargs) -> np.ndarray:
+        all_embeddings = []
+        for i in range(0, len(sentences), batch_size):
+            batch = sentences[i : i + batch_size]
+            response = self._model.create_embedding(batch)
+            embeddings = [item["embedding"] for item in response["data"]]
+            all_embeddings.extend(embeddings)
+        return np.array(all_embeddings, dtype=np.float32)
+
+
+class FastEmbedWrapper:
+    """Wraps fastembed.TextEmbedding."""
+
+    def __init__(self, cfg: ModelConfig):
+        from fastembed import TextEmbedding
+        self._model = TextEmbedding(model_name=cfg.model_id)
+
+    def encode(self, sentences: list[str], batch_size: int = 64, **kwargs) -> np.ndarray:
+        embeddings = list(self._model.embed(sentences, batch_size=batch_size))
+        return np.array(embeddings, dtype=np.float32)
+
+
+def load_model(cfg: ModelConfig) -> SBertWrapper | GGUFWrapper | FastEmbedWrapper:
+    """Factory: returns the right wrapper for the model's backend."""
+    if cfg.backend == "gguf":
+        return GGUFWrapper(cfg)
+    if cfg.backend == "fastembed":
+        return FastEmbedWrapper(cfg)
+    return SBertWrapper(cfg)