Embedding

embedding

Embedding backends + brute-force cosine neighbor search.

Two registered backends, ported from the dedup-detection POC (find_duplicates_v7.py + trials/trial_unified.py):

• tfidf : TF-IDF over code tokens (scikit-learn). Pure CPU, no torch.
• st : sentence-transformers MiniLM (all-MiniLM-L6-v2).

torch / sentence_transformers are base dependencies (AXM-2188): the st backend is the default and runs in-process. The heavy model import is still done lazily inside the st backend purely for first-call latency, so the tfidf path stays light and never loads torch.

Neighbor search is an exact brute-force matmul (no ANN): under ~10^5 vectors a single normalized dot product beats an approximate index by 7-30x while staying exact. This calibration is closed (see ticket technical notes) and is not re-litigated here.

code_tokens(src)

Tokenize code text for the TF-IDF backend.

Lowercases identifiers, splits camelCase and snake_case into sub-tokens, and appends structural keyword hints weighted by frequency so control-flow shape contributes to the vector.

Source code in packages/axm-echo/src/axm_echo/embedding.py

def code_tokens(src: str) -> list[str]:
    """Tokenize code text for the TF-IDF backend.

    Lowercases identifiers, splits ``camelCase`` and ``snake_case`` into
    sub-tokens, and appends structural keyword hints weighted by frequency
    so control-flow shape contributes to the vector.
    """
    tokens: list[str] = []
    for ident in _IDENT_RE.findall(src):
        tokens.append(ident.lower())
        parts = re.split(r"(?<=[a-z])(?=[A-Z])|_", ident)
        tokens.extend(p.lower() for p in parts if p)
    ident_tokens = [t for t in tokens if t in _STRUCTURAL_KEYWORDS]
    for kw in _STRUCTURAL_KEYWORDS:
        tokens.extend([kw] * ident_tokens.count(kw))
    return tokens

embed(texts, *, backend='tfidf')

Embed texts into a dense matrix via the chosen backend.

Parameters:

Name	Type	Description	Default
texts	Sequence[str]	Texts to embed (one row per text).	required
backend	Backend	"tfidf" (code, scikit-learn) or "st" (MiniLM neural). The st backend is the tools' default (this lib-layer function defaults to tfidf for a torch-free call); it lazily imports torch for first-call latency, tfidf never loads it.	'tfidf'

Returns:

Type	Description
NDArray[float64]	A (len(texts), dim) float matrix. Rows are not guaranteed
NDArray[float64]	normalized; neighbors normalizes internally.

Raises:

Type	Description
ValueError	If backend is not a registered backend.

Source code in packages/axm-echo/src/axm_echo/embedding.py

def embed(texts: Sequence[str], *, backend: Backend = "tfidf") -> NDArray[np.float64]:
    """Embed ``texts`` into a dense matrix via the chosen backend.

    Args:
        texts: Texts to embed (one row per text).
        backend: ``"tfidf"`` (code, scikit-learn) or ``"st"`` (MiniLM neural).
            The ``st`` backend is the tools' default (this lib-layer function
            defaults to ``tfidf`` for a torch-free call); it lazily imports
            torch for first-call latency, ``tfidf`` never loads it.

    Returns:
        A ``(len(texts), dim)`` float matrix. Rows are not guaranteed
        normalized; ``neighbors`` normalizes internally.

    Raises:
        ValueError: If ``backend`` is not a registered backend.
    """
    try:
        fn = _BACKENDS[backend]
    except KeyError:
        valid = ", ".join(sorted(_BACKENDS))
        msg = f"unknown backend {backend!r}; expected one of: {valid}"
        raise ValueError(msg) from None
    return fn(texts)

neighbors(query, matrix, *, k=10, threshold=None)

Cosine top-k nearest rows to query (exact brute-force matmul).

Parameters:

Name	Type	Description	Default
query	NDArray[float64]	A single query vector of shape (dim,).	required
matrix	NDArray[float64]	Corpus matrix of shape (n, dim).	required
k	int	Maximum number of neighbors to return.	10
threshold	float | None	If set, drop neighbors with cosine below this value.	None

Returns:

Type	Description
list[tuple[int, float]]	A list of (row_index, cosine_score) pairs sorted by score
list[tuple[int, float]]	descending, length <= k. k <= 0 yields an empty list. The
list[tuple[int, float]]	query is not excluded from the corpus; if query is a row of
list[tuple[int, float]]	matrix it appears at score 1.0.

Source code in packages/axm-echo/src/axm_echo/embedding.py

def neighbors(
    query: NDArray[np.float64],
    matrix: NDArray[np.float64],
    *,
    k: int = 10,
    threshold: float | None = None,
) -> list[tuple[int, float]]:
    """Cosine top-k nearest rows to ``query`` (exact brute-force matmul).

    Args:
        query: A single query vector of shape ``(dim,)``.
        matrix: Corpus matrix of shape ``(n, dim)``.
        k: Maximum number of neighbors to return.
        threshold: If set, drop neighbors with cosine below this value.

    Returns:
        A list of ``(row_index, cosine_score)`` pairs sorted by score
        descending, length ``<= k``. ``k <= 0`` yields an empty list. The
        query is not excluded from the corpus; if ``query`` is a row of
        ``matrix`` it appears at score 1.0.
    """
    if k <= 0:
        return []
    norm_matrix = _l2_normalize(np.atleast_2d(matrix).astype(np.float64))
    norm_query = _l2_normalize(np.asarray(query, dtype=np.float64).reshape(1, -1))
    sims = (norm_query @ norm_matrix.T).ravel()

    order = np.argsort(-sims)
    results: list[tuple[int, float]] = []
    for idx in order:
        score = float(sims[idx])
        if threshold is not None and score < threshold:
            break
        results.append((int(idx), score))
        if len(results) >= k:
            break
    return results