Spaces:
Running
Running
File size: 3,531 Bytes
fb0ec21 | 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 | """Embedding providers for the ABM vector memory.
`auto` prefers fastembed (ONNX MiniLM, 384-dim, real semantic similarity) and
falls back to a deterministic hashing embedding when the model is unavailable
(cold boot, build without the model cache). Both produce L2-normalized vectors
of settings.embedding_dim so the k-NN index stays compatible either way.
"""
import hashlib
import logging
import math
import re
from typing import List
from ..config import settings
logger = logging.getLogger("abm.embeddings")
_TOKEN_RE = re.compile(r"[a-z0-9]{2,}")
class HashEmbedding:
"""Deterministic feature-hashing embedding (bag of tokens + bigrams).
Honest about what it is: lexical-ish similarity, not deep semantics.
Zero dependencies, instant, stable across restarts — the safe default.
"""
name = "hash-v1"
def _dims(self) -> int:
return settings.embedding_dim
def _index(self, token: str) -> tuple[int, float]:
digest = hashlib.md5(token.encode("utf-8")).digest()
idx = int.from_bytes(digest[:4], "little") % self._dims()
sign = 1.0 if digest[4] % 2 == 0 else -1.0
return idx, sign
def encode(self, text: str) -> List[float]:
vec = [0.0] * self._dims()
tokens = _TOKEN_RE.findall((text or "").lower())
grams: List[str] = list(tokens)
grams += [f"{a}_{b}" for a, b in zip(tokens, tokens[1:])]
for g in grams:
idx, sign = self._index(g)
vec[idx] += sign
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
return [round(v / norm, 6) for v in vec]
def encode_batch(self, texts: List[str]) -> List[List[float]]:
return [self.encode(t) for t in texts]
class FastEmbedEmbedding:
"""ONNX MiniLM via fastembed — real semantic similarity, no torch."""
def __init__(self) -> None:
from fastembed import TextEmbedding # imported lazily
self._model = TextEmbedding("sentence-transformers/all-MiniLM-L6-v2")
self.name = "fastembed-minilm-384"
def _normalize(self, vec) -> List[float]:
norm = math.sqrt(sum(v * v for v in vec)) or 1.0
return [round(float(v) / norm, 6) for v in vec]
def encode(self, text: str) -> List[float]:
out = list(self._model.embed([text]))[0]
return self._normalize(out)
def encode_batch(self, texts: List[str]) -> List[List[float]]:
return [self._normalize(v) for v in self._model.embed(texts)]
_provider = None
_tried = False
def get_provider():
global _provider, _tried
if _provider is not None:
return _provider
if _tried:
return None
_tried = True
choice = (settings.embedding_provider or "auto").lower()
if choice in ("auto", "fastembed"):
try:
_provider = FastEmbedEmbedding()
logger.info("embeddings: fastembed active (%s)", _provider.name)
return _provider
except Exception as exc:
logger.warning("fastembed unavailable (%s) — using hash embeddings", exc)
_provider = HashEmbedding()
return _provider
def embed(text: str) -> List[float]:
p = get_provider()
if p is None:
return HashEmbedding().encode(text)
return p.encode(text)
def embed_batch(texts: List[str]) -> List[List[float]]:
p = get_provider()
if p is None:
return HashEmbedding().encode_batch(texts)
return p.encode_batch(texts)
def provider_name() -> str:
p = get_provider()
return p.name if p else "hash-v1"
|