Spaces:
Sleeping
Sleeping
File size: 4,983 Bytes
454f118 | 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 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """
src/embeddings/embedder.py
===========================
Embedding engine using BAAI/bge-large-en-v1.5.
Why BGE-large?
- Free, runs locally on your RTX 3050
- Top of MTEB leaderboard for retrieval tasks (among free models)
- 768-dimensional output — good balance of quality vs. memory
- Trained with specific instruction prefix for queries vs. documents
BGE-specific usage:
Documents → embed as-is (no prefix)
Queries → add "Represent this sentence: " prefix
This asymmetric approach improves retrieval accuracy significantly.
L2 Normalization:
All embeddings are L2-normalized before storage.
This makes cosine similarity = dot product → faster ANN search in Qdrant.
"""
from __future__ import annotations
import numpy as np
from typing import Union
from tqdm import tqdm
# config import first — sets USE_TF=0 before sentence_transformers loads
from config.settings import settings
from src.utils.logger import logger
# Now safe to import
from sentence_transformers import SentenceTransformer
class Embedder:
"""
Singleton-pattern embedding engine.
Loads BAAI/bge-large-en-v1.5 once and reuses it.
GPU-accelerated on RTX 3050 automatically (via settings.get_device()).
"""
_instance: "Embedder | None" = None
def __new__(cls):
"""Enforce singleton — only load the model once."""
if cls._instance is None:
cls._instance = super().__new__(cls)
cls._instance._initialized = False
return cls._instance
def __init__(self):
if self._initialized:
return
self._load_model()
self._initialized = True
def _load_model(self):
"""Load the embedding model onto GPU/CPU."""
model_name = settings.embedding_model
device = settings.get_device()
logger.info(f"Loading embedding model: {model_name} → device: {device}")
logger.info("(First load downloads ~1.3GB — subsequent loads use cache)")
self.model = SentenceTransformer(model_name, device=device)
self.model.max_seq_length = 512 # BGE-large max input length
self.dim = settings.embedding_dim # 768
logger.info(f"Embedding model loaded | dim={self.dim} | device={device}")
# ---------------------------------------------------------------- #
# Public API
# ---------------------------------------------------------------- #
def embed_documents(self, texts: list[str]) -> np.ndarray:
"""
Embed a list of document chunks (corpus side).
No instruction prefix needed for documents with BGE.
Args:
texts: List of text strings to embed
Returns:
np.ndarray of shape (len(texts), 768), L2-normalized
"""
if not texts:
return np.array([])
logger.info(f"Embedding {len(texts)} documents in batches of {settings.embed_batch_size}...")
embeddings = []
batch_size = settings.embed_batch_size
for i in tqdm(range(0, len(texts), batch_size), desc="Embedding", unit="batch"):
batch = texts[i: i + batch_size]
batch_emb = self.model.encode(
batch,
batch_size=batch_size,
show_progress_bar=False,
normalize_embeddings=True, # L2 normalize in-place
convert_to_numpy=True,
)
embeddings.append(batch_emb)
result = np.vstack(embeddings)
logger.info(f"Embedding complete: shape={result.shape}")
return result
def embed_query(self, query: str) -> np.ndarray:
"""
Embed a single search query (query side).
BGE models benefit from an instruction prefix on queries.
Args:
query: The user's question
Returns:
np.ndarray of shape (768,), L2-normalized
"""
# BGE instruction prefix — improves retrieval accuracy
prefixed_query = f"Represent this sentence for searching relevant passages: {query}"
embedding = self.model.encode(
prefixed_query,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=False,
)
return embedding
def embed_queries(self, queries: list[str]) -> np.ndarray:
"""
Embed multiple queries (for multi-query retrieval expansion).
Args:
queries: List of question strings
Returns:
np.ndarray of shape (len(queries), 768)
"""
prefixed = [
f"Represent this sentence for searching relevant passages: {q}"
for q in queries
]
embeddings = self.model.encode(
prefixed,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=False,
)
return embeddings
@property
def embedding_dim(self) -> int:
return self.dim
|