Add vrom_hub/embedder.py

vrom_hub/embedder.py

@@ -0,0 +1,89 @@
+"""
+Chunk embedder using sentence-transformers.
+
+Produces 384-dimensional normalized embeddings compatible with
+the vROM ecosystem (Xenova/all-MiniLM-L6-v2, cosine metric).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from vrom_hub.chunker import Chunk
+
+logger = logging.getLogger(__name__)
+
+
+class ChunkEmbedder:
+    """
+    Embeds chunk text using sentence-transformers/all-MiniLM-L6-v2.
+
+    The model produces 384-dimensional embeddings. Vectors are L2-normalized
+    for cosine similarity (consistent with the WASM runtime).
+    """
+
+    def __init__(
+        self,
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        device: str | None = None,
+        batch_size: int = 64,
+    ):
+        self.model_name = model_name
+        self.batch_size = batch_size
+        self._model = None
+        self._device = device
+
+    @property
+    def model(self):
+        if self._model is None:
+            from sentence_transformers import SentenceTransformer
+            self._model = SentenceTransformer(self.model_name, device=self._device)
+            logger.info(
+                f"Loaded embedding model: {self.model_name} "
+                f"(dim={self._model.get_embedding_dimension()})"
+            )
+        return self._model
+
+    @property
+    def dimensions(self) -> int:
+        return self.model.get_embedding_dimension()
+
+    def embed_texts(self, texts: list[str]) -> np.ndarray:
+        """
+        Embed a list of texts.
+
+        Returns:
+            np.ndarray of shape (len(texts), dim) with L2-normalized vectors.
+        """
+        logger.info(f"Embedding {len(texts)} texts in batches of {self.batch_size}...")
+        embeddings = self.model.encode(
+            texts,
+            batch_size=self.batch_size,
+            show_progress_bar=True,
+            normalize_embeddings=True,  # L2-normalize for cosine
+            convert_to_numpy=True,
+        )
+        logger.info(f"Embeddings shape: {embeddings.shape}")
+        return embeddings
+
+    def embed_chunks(self, chunks: list[Chunk]) -> np.ndarray:
+        """
+        Embed a list of Chunk objects by their text content.
+
+        Returns:
+            np.ndarray of shape (len(chunks), dim) with L2-normalized vectors.
+        """
+        texts = [c.text for c in chunks]
+        return self.embed_texts(texts)
+
+    def embed_query(self, query: str) -> np.ndarray:
+        """Embed a single query string. Returns shape (dim,)."""
+        return self.model.encode(
+            [query],
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+        )[0]