File size: 4,980 Bytes
722bda8 | 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 | """
Text embedding wrapper for scene descriptions.
Uses BAAI/bge-m3 via sentence-transformers. bge-m3 is multilingual (JW.org
content exists in many languages) and produces 1024-dim dense vectors that
slot into the same sqlite-vec storage pattern the rest of the project uses.
"""
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import Sequence
from utils import log_message
DEFAULT_MODEL_ID = "BAAI/bge-m3"
DEFAULT_BATCH_SIZE = 16
# bge-m3 has an 8192-token context. ~3.5 chars/token average for our
# descriptions, so 24000 chars is a safe per-chunk budget that leaves
# headroom for tokenization edge cases.
SUMMARY_CHUNK_CHAR_BUDGET = 24000
@dataclass
class EmbeddingResult:
"""Embedding for one scene description."""
vector: list[float]
model_id: str
dimension: int
class TextEmbedder:
"""Lazy-loaded sentence-transformers wrapper."""
def __init__(
self,
model_id: str = DEFAULT_MODEL_ID,
*,
batch_size: int = DEFAULT_BATCH_SIZE,
) -> None:
self.model_id = model_id
self.batch_size = batch_size
self._model = None
def _ensure_loaded(self) -> None:
if self._model is not None:
return
try:
from sentence_transformers import SentenceTransformer
except ImportError as exc:
raise RuntimeError(
"sentence-transformers is required for the text embedder."
) from exc
log_message(f"Loading text embedder: {self.model_id}")
self._model = SentenceTransformer(self.model_id)
log_message(f"Text embedder ready: {self.model_id}")
def embed_many(self, texts: Sequence[str]) -> list[EmbeddingResult]:
"""Embed a list of texts. Returns one ``EmbeddingResult`` per input."""
self._ensure_loaded()
if not texts:
return []
vectors = self._model.encode(
list(texts),
batch_size=self.batch_size,
normalize_embeddings=True,
convert_to_numpy=True,
show_progress_bar=False,
)
return [
EmbeddingResult(
vector=vector.tolist(),
model_id=self.model_id,
dimension=int(vector.shape[0]),
)
for vector in vectors
]
def chunk_pieces_for_embedding(
pieces: Sequence[str],
*,
char_budget: int = SUMMARY_CHUNK_CHAR_BUDGET,
) -> list[str]:
"""Group text pieces into newline-joined chunks under a character budget.
Used to keep summary embeddings from silently truncating when the joined
text exceeds the embedder's context window. Each output chunk is the
newline-joined text of a contiguous slice of ``pieces`` whose total
length stays under ``char_budget``.
"""
chunks: list[str] = []
current: list[str] = []
current_len = 0
for piece in pieces:
piece_len = len(piece) + 1 # +1 for the newline that will join them
if current and current_len + piece_len > char_budget:
chunks.append("\n".join(current))
current = [piece]
current_len = piece_len
else:
current.append(piece)
current_len += piece_len
if current:
chunks.append("\n".join(current))
return chunks
def mean_pool_normalized(vectors: Sequence[Sequence[float]]) -> list[float]:
"""Mean across vectors then re-normalize to unit length.
Used to roll multiple chunk embeddings into one representative vector
that still lives on the unit sphere (so cosine similarity remains
well-defined against single-chunk embeddings).
"""
if not vectors:
return []
dim = len(vectors[0])
pooled = [0.0] * dim
for vec in vectors:
for i in range(dim):
pooled[i] += vec[i]
n = len(vectors)
pooled = [x / n for x in pooled]
norm = math.sqrt(sum(x * x for x in pooled))
if norm == 0:
return pooled
return [x / norm for x in pooled]
def embed_pieces_pooled(
pieces: Sequence[str],
embedder: TextEmbedder,
) -> tuple[list[float], int, str]:
"""Chunk ``pieces`` to fit the embedder, embed each chunk, mean-pool the result.
Returns ``(vector, dimension, model_id)``. For an empty input returns
``([], 0, embedder.model_id)`` so callers can detect "nothing to embed"
without checking length on the vector themselves.
"""
if not pieces:
return [], 0, embedder.model_id
chunks = chunk_pieces_for_embedding(pieces)
chunk_embeddings = embedder.embed_many(chunks)
if not chunk_embeddings:
return [], 0, embedder.model_id
if len(chunk_embeddings) == 1:
emb = chunk_embeddings[0]
return emb.vector, emb.dimension, emb.model_id
pooled = mean_pool_normalized([e.vector for e in chunk_embeddings])
return pooled, chunk_embeddings[0].dimension, chunk_embeddings[0].model_id
|