File size: 4,160 Bytes
939c0c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | """
Embedding Service
=================
Generates dense vector embeddings using sentence-transformers (100% free, runs locally).
Model: all-MiniLM-L6-v2 (384-dim, fast, great quality for enterprise text)
Singleton pattern: model loads once on first call and stays in memory.
Thread-safe for concurrent FastAPI requests.
"""
from __future__ import annotations
import asyncio
import threading
from typing import List, Optional
import numpy as np
import structlog
logger = structlog.get_logger(__name__)
class EmbeddingService:
"""
Wraps sentence-transformers SentenceTransformer for async-compatible embedding.
Falls back gracefully if the model is not yet downloaded (first run).
"""
def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
self.model_name = model_name
self._model = None
self._lock = threading.Lock()
self._dim: Optional[int] = None
def _load_model(self):
"""Lazy-load the embedding model (downloads on first use)."""
if self._model is None:
with self._lock:
if self._model is None:
logger.info(
"Loading embedding model",
model=self.model_name,
note="Downloading on first run, ~90MB",
)
from sentence_transformers import SentenceTransformer
self._model = SentenceTransformer(self.model_name)
# Warm up with a test sentence
test = self._model.encode(["test"], show_progress_bar=False)
self._dim = test.shape[1]
logger.info(
"Embedding model ready",
model=self.model_name,
dim=self._dim,
)
@property
def dim(self) -> int:
"""Return embedding dimensions (loaded lazily)."""
if self._dim is None:
self._load_model()
return self._dim # type: ignore
def encode_sync(
self,
texts: List[str],
batch_size: int = 64,
normalize: bool = True,
) -> np.ndarray:
"""
Synchronous encoding — for use in background tasks.
Args:
texts: List of strings to embed
batch_size: Batch size for encoding (larger = faster on GPU)
normalize: L2-normalize embeddings (recommended for cosine similarity)
Returns:
numpy array of shape (len(texts), dim)
"""
self._load_model()
embeddings = self._model.encode( # type: ignore
texts,
batch_size=batch_size,
normalize_embeddings=normalize,
show_progress_bar=False,
)
return embeddings
async def encode(
self,
texts: List[str],
batch_size: int = 64,
normalize: bool = True,
) -> np.ndarray:
"""
Async encoding — runs encoding in a thread pool to avoid blocking the event loop.
Args:
texts: List of strings to embed
batch_size: Batch size for encoding
normalize: L2-normalize embeddings
Returns:
numpy array of shape (len(texts), dim)
"""
loop = asyncio.get_event_loop()
return await loop.run_in_executor(
None, lambda: self.encode_sync(texts, batch_size, normalize)
)
async def encode_query(self, query: str) -> np.ndarray:
"""Encode a single query string and return 1D array."""
result = await self.encode([query])
return result[0]
async def encode_documents(self, texts: List[str]) -> np.ndarray:
"""
Encode a list of document chunks. Adds a prefix per MTEB best practice.
"""
prefixed = [f"passage: {t}" for t in texts]
return await self.encode(prefixed)
# ── Singleton instance ─────────────────────────────────────────
embedding_service = EmbeddingService(model_name="all-MiniLM-L6-v2")
|