egutierrez
0fa16a033c
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>
52 lines
1.7 KiB
Python
52 lines
1.7 KiB
Python
"""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))]
|