Embedding Model

EmbeddingModel

EmbeddingModel(
    model_name="all-MiniLM-L6-v2",
    device=None,
    cache_embeddings=True,
)

Wrapper around SentenceTransformer with lazy loading and caching.

Provides thread-safe lazy loading of embedding models and optional caching of computed embeddings for improved performance. Supports both single text encoding and batch processing.

Attributes:

Name	Type	Description
model_name		Name of the sentence transformer model to use.
device		Device to run the model on ("cuda", "cpu", or None for auto).
cache_embeddings		Whether to cache computed embeddings using LRU cache.

Initialize the embedding model wrapper.

Source code in at_scorer/ml/embeddings.py

def __init__(
    self,
    model_name: str = "all-MiniLM-L6-v2",
    device: str | None = None,
    cache_embeddings: bool = True,
):
    """Initialize the embedding model wrapper."""
    self.model_name = model_name
    self.device = device
    self.cache_embeddings = cache_embeddings

Functions

encode

encode(text)

Encode a single text string into an embedding vector.

Parameters:

Name	Type	Description	Default
text	str	The text to encode.	required

Returns:

Type	Description
ndarray	Normalized embedding vector as numpy array. Returns zero vector for empty text.

Source code in at_scorer/ml/embeddings.py

def encode(self, text: str) -> np.ndarray:
    """Encode a single text string into an embedding vector.

    Args:
        text: The text to encode.

    Returns:
        Normalized embedding vector as numpy array. Returns zero vector for empty text.
    """
    if not text:
        return np.zeros(self._vector_size(), dtype=float)
    if self.cache_embeddings:
        return self._encode_cached(text)
    return self._encode_uncached(text)

encode_batch

encode_batch(texts)

Encode multiple texts efficiently in a batch.

Parameters:

Name	Type	Description	Default
texts	Iterable[str]	Iterable of text strings to encode.	required

Returns:

Type	Description
list[ndarray]	List of normalized embedding vectors as numpy arrays.

Source code in at_scorer/ml/embeddings.py

def encode_batch(self, texts: Iterable[str]) -> list[np.ndarray]:
    """Encode multiple texts efficiently in a batch.

    Args:
        texts: Iterable of text strings to encode.

    Returns:
        List of normalized embedding vectors as numpy arrays.
    """
    texts_list = list(texts)
    if not texts_list:
        return []
    model = self._get_model()
    embeddings = model.encode(texts_list, normalize_embeddings=True)
    return [np.array(vec) for vec in embeddings]

similarity

similarity(a, b)

Calculate cosine similarity between two embedding vectors.

Parameters:

Name	Type	Description	Default
a	ndarray	First embedding vector.	required
b	ndarray	Second embedding vector.	required

Returns:

Type	Description
float	Cosine similarity score between -1 and 1 (typically 0-1 for normalized vectors).
float	Returns 0.0 if either vector is empty.

Source code in at_scorer/ml/embeddings.py

def similarity(self, a: np.ndarray, b: np.ndarray) -> float:
    """Calculate cosine similarity between two embedding vectors.

    Args:
        a: First embedding vector.
        b: Second embedding vector.

    Returns:
        Cosine similarity score between -1 and 1 (typically 0-1 for normalized vectors).
        Returns 0.0 if either vector is empty.
    """
    if a.size == 0 or b.size == 0:
        return 0.0
    denom = (np.linalg.norm(a) * np.linalg.norm(b)) or 1e-9
    return float(np.dot(a, b) / denom)