feat: revert backend to use sentence-transformers (all-MiniLM-L6-v2) for better embeddings

backend/api.py

@@ -5,14 +5,15 @@ from typing import Any, Dict, List
 
 import numpy as np
 import spacy
+from spacy import displacy
 from fastapi import FastAPI, Query
 from fastapi.middleware.cors import CORSMiddleware
 from rank_bm25 import BM25Okapi
 
 try:
-    from fastembed import TextEmbedding
+    from sentence_transformers import SentenceTransformer
 except Exception:
-    TextEmbedding = None
+    SentenceTransformer = None
 
 
 BASE_DIR = Path(__file__).resolve().parents[1]
@@ -95,11 +96,10 @@ def _role_overlap_score(
 
 
 def _maybe_load_embedder() -> Any:
-    if TextEmbedding is None:
+    if SentenceTransformer is None:
         return None
     try:
-        # threads=1 prevents FastEmbed from spawning RAM-hungry multiprocessing workers
-        return TextEmbedding("sentence-transformers/all-MiniLM-L6-v2", threads=1)
+        return SentenceTransformer("all-MiniLM-L6-v2")
     except Exception as e:
         print(f"Failed to load embedder: {e}")
         return None
@@ -109,8 +109,7 @@ def _dense_scores(query: str) -> np.ndarray:
     if EMBEDDER is None or DOC_EMBEDDINGS.size == 0:
         return np.zeros(len(CORPUS), dtype=float)
 
-    # FastEmbed returns a generator of arrays, we need the first (and only) one
-    query_embedding = list(EMBEDDER.embed([query]))[0]
+    query_embedding = EMBEDDER.encode(query)
     return np.array(np.dot(DOC_EMBEDDINGS, query_embedding), dtype=float)
 
 
@@ -145,7 +144,7 @@ def _build_result(
 
     try:
         spacy_doc = NLP(text)
-        html = spacy.displacy.render(
+        html = displacy.render(
             spacy_doc,
             style="dep",
             page=False,
@@ -192,7 +191,7 @@ if EMBEDDINGS_PATH.exists():
         DOC_EMBEDDINGS = DOC_EMBEDDINGS[:MAX_DOCS]
 elif EMBEDDER is not None and TEXTS:
     print("Computing embeddings on the fly... (This may cause OOM on small instances)")
-    DOC_EMBEDDINGS = np.array(list(EMBEDDER.embed(TEXTS)))
+    DOC_EMBEDDINGS = np.array(EMBEDDER.encode(TEXTS))
 else:
     DOC_EMBEDDINGS = np.array([])
 

backend/requirements.txt

@@ -1,6 +1,6 @@
 fastapi
 uvicorn
 rank-bm25
-fastembed
+sentence-transformers
 spacy
 numpy