feat: módulo embedding — encode, model CRUD, stores sqlvec y usearch
Funciones Python para embeddings: carga/guardado de modelos, encoding de texto, y almacenamiento/búsqueda vectorial con sqlite-vec y usearch. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,51 @@
|
||||
"""Vector storage and retrieval with USearch."""
|
||||
|
||||
import numpy as np
|
||||
from usearch.index import Index
|
||||
|
||||
|
||||
def embedding_store_usearch(path: str, ids: list, embeddings: list, dim: int = 384) -> int:
|
||||
"""Crea o sobreescribe un indice USearch con los embeddings dados.
|
||||
|
||||
Args:
|
||||
path: Path del archivo .usearch para persistir el indice.
|
||||
ids: Lista de IDs (int) para cada embedding.
|
||||
embeddings: Lista de arrays numpy float32.
|
||||
dim: Dimension de los embeddings (default 384 para e5-small).
|
||||
|
||||
Returns:
|
||||
Numero de embeddings insertados.
|
||||
|
||||
Raises:
|
||||
OSError: Si no se puede escribir en el path.
|
||||
"""
|
||||
index = Index(ndim=dim, metric="ip", dtype="f32")
|
||||
keys = np.array(ids, dtype=np.uint64)
|
||||
vecs = np.array(embeddings, dtype=np.float32)
|
||||
index.add(keys, vecs)
|
||||
index.save(path)
|
||||
return len(ids)
|
||||
|
||||
|
||||
def embedding_search_usearch(path: str, query_embedding: list, k: int = 10, dim: int = 384) -> list:
|
||||
"""Busca los k vecinos mas cercanos en un indice USearch persistido.
|
||||
|
||||
Args:
|
||||
path: Path del archivo .usearch.
|
||||
query_embedding: Array numpy float32 del query.
|
||||
k: Numero de resultados a retornar.
|
||||
dim: Dimension de los embeddings.
|
||||
|
||||
Returns:
|
||||
Lista de dicts con 'key' y 'distance' ordenados por cercania.
|
||||
|
||||
Raises:
|
||||
OSError: Si el archivo no existe.
|
||||
"""
|
||||
index = Index(ndim=dim, metric="ip", dtype="f32")
|
||||
index.load(path)
|
||||
q = np.asarray(query_embedding, dtype=np.float32)
|
||||
results = index.search(q, k)
|
||||
keys = np.atleast_1d(results.keys)
|
||||
distances = np.atleast_1d(results.distances)
|
||||
return [{"key": int(keys[i]), "distance": float(distances[i])} for i in range(len(keys))]
|
||||
Reference in New Issue
Block a user