File size: 945 Bytes
d10de1b
 
 
 
 
e9364cc
d10de1b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Encode text chunks into dense vectors using SentenceTransformers."""

import numpy as np
from sentence_transformers import SentenceTransformer

from config import EMBEDDING_MODEL

_model = None  # lazy singleton


def get_embedding_model() -> SentenceTransformer:
    global _model
    if _model is None:
        print(f"[embeddings] Loading embedding model: {EMBEDDING_MODEL}")
        _model = SentenceTransformer(EMBEDDING_MODEL)
    return _model


def embed_texts(texts: list[str]) -> np.ndarray:
    """Return a (N, D) float32 array of embeddings."""
    model = get_embedding_model()
    embeddings = model.encode(texts, show_progress_bar=True, convert_to_numpy=True)
    return embeddings.astype("float32")


def embed_query(query: str) -> np.ndarray:
    """Return a (1, D) float32 array for a single query."""
    model = get_embedding_model()
    vec = model.encode([query], convert_to_numpy=True)
    return vec.astype("float32")