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,67 @@
|
||||
"""Embedding model management — save, load, and encode with multilingual-e5-small."""
|
||||
|
||||
import os
|
||||
|
||||
from sentence_transformers import SentenceTransformer
|
||||
|
||||
|
||||
def embedding_save_model(model_id: str, path: str) -> str:
|
||||
"""Descarga modelo de HuggingFace y lo guarda en path local.
|
||||
|
||||
Args:
|
||||
model_id: ID del modelo en HuggingFace (ej: "intfloat/multilingual-e5-small").
|
||||
path: Directorio destino para guardar el modelo.
|
||||
|
||||
Returns:
|
||||
Path absoluto donde se guardo el modelo.
|
||||
|
||||
Raises:
|
||||
OSError: Si no se puede escribir en el path.
|
||||
Exception: Si el modelo no existe en HuggingFace.
|
||||
"""
|
||||
os.makedirs(path, exist_ok=True)
|
||||
model = SentenceTransformer(model_id)
|
||||
model.save(path)
|
||||
return os.path.abspath(path)
|
||||
|
||||
|
||||
def embedding_load_model(path: str) -> SentenceTransformer:
|
||||
"""Carga modelo de embeddings desde path local.
|
||||
|
||||
Args:
|
||||
path: Directorio con el modelo guardado por embedding_save_model.
|
||||
|
||||
Returns:
|
||||
Instancia de SentenceTransformer lista para encode.
|
||||
|
||||
Raises:
|
||||
OSError: Si el path no existe o no contiene un modelo valido.
|
||||
"""
|
||||
return SentenceTransformer(path)
|
||||
|
||||
|
||||
def embedding_encode(model: SentenceTransformer, texts: list, mode: str = "document") -> list:
|
||||
"""Genera embeddings normalizados para una lista de textos.
|
||||
|
||||
Aplica automaticamente los prefijos requeridos por modelos e5:
|
||||
- mode="document" -> "passage: " prefix
|
||||
- mode="query" -> "query: " prefix
|
||||
|
||||
Args:
|
||||
model: Modelo cargado con embedding_load_model.
|
||||
texts: Lista de strings a codificar.
|
||||
mode: "document" para indexar, "query" para buscar.
|
||||
|
||||
Returns:
|
||||
Lista de arrays numpy float32 normalizados (dim depende del modelo).
|
||||
|
||||
Raises:
|
||||
ValueError: Si mode no es "document" ni "query".
|
||||
"""
|
||||
if mode not in ("document", "query"):
|
||||
raise ValueError(f"mode must be 'document' or 'query', got '{mode}'")
|
||||
|
||||
prefix = "passage: " if mode == "document" else "query: "
|
||||
prefixed = [f"{prefix}{t}" for t in texts]
|
||||
embeddings = model.encode(prefixed, normalize_embeddings=True, show_progress_bar=False)
|
||||
return embeddings
|
||||
Reference in New Issue
Block a user