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>
68 lines
2.1 KiB
Python
68 lines
2.1 KiB
Python
"""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
|