La mejor forma de usar embeddings de texto de manera portable es con Parquet y Polars
- Los embeddings de texto son vectores generados por modelos de lenguaje grandes, una forma de representar numéricamente palabras, oraciones y documentos
- A febrero de 2025, se generaron embeddings para un total de 32,254 cartas de "Magic: The Gathering"
- Esto permite analizar matemáticamente la similitud entre cartas con base en sus propiedades de diseño y mecánicas
- Los embeddings generados pueden visualizarse mediante reducción de dimensionalidad 2D con UMAP
- El modelo de embeddings utilizado es gte-modernbert-base, y el proceso detallado está documentado en el repositorio de GitHub
- Este dataset de embeddings está disponible en Hugging Face
Replanteando la necesidad de una base de datos vectorial
- Normalmente se usan bases de datos vectoriales (faiss, qdrant, Pinecone) para almacenar y buscar embeddings
- Sin embargo, las bases de datos vectoriales requieren configuraciones complejas, y los servicios en la nube pueden ser costosos
- Para datos de pequeña escala (del orden de decenas de miles), es posible hacer búsquedas de similitud rápidas con numpy incluso sin una base de datos vectorial
- Aprovechando la operación de
dot product de numpy, se puede calcular similitud coseno de forma simple, y para 32,254 embeddings toma en promedio 1.08 ms
def fast_dot_product(query, matrix, k=3):
dot_products = query @ matrix.T
idx = np.argpartition(dot_products, -k)[-k:]
idx = idx[np.argsort(dot_products[idx])[::-1]]
score = dot_products[idx]
return idx, score
- Al usar una base de datos vectorial, hay una alta probabilidad de quedar atado a bibliotecas y servicios específicos
- Si se generan embeddings en un servidor con GPU y luego se descargan localmente, se necesita una forma eficiente de almacenar y transferir los datos
La peor forma de guardar embeddings
- Archivos CSV
- Si se almacenan datos de punto flotante (
float32) como texto, el tamaño aumenta más de 6 veces
- Incluso el tutorial oficial de OpenAI recomienda usar CSV solo para datasets pequeños
- Guardarlos con
.savetxt() de numpy hace que el archivo crezca hasta 631.5MB
- Archivos pickle
- Permiten guardar y cargar rápido, pero tienen riesgos de seguridad y poca compatibilidad entre versiones
- El tamaño del archivo es 94.49MB, igual al tamaño original en memoria, pero con baja portabilidad
Formas de almacenamiento que no están mal, pero no son óptimas
- Formato
.npy de numpy
- Con
allow_pickle=False se puede evitar el almacenamiento con pickle
- El tamaño del archivo y la velocidad son iguales al método pickle, pero es difícil guardar metadatos individuales junto con los embeddings
- Problemas de una estructura de almacenamiento separada de los metadatos
- Si se guarda como un arreglo numpy (
.npy), la información de las cartas (nombre, texto, etc.) queda separada de los embeddings
- Cuando los datos cambian (altas o bajas), se vuelve difícil mantener la correspondencia entre metadatos y embeddings
- En una base de datos vectorial, los metadatos y los vectores se guardan juntos y se ofrecen funciones de filtrado
La mejor forma de guardar embeddings: Parquet + polars
Introducción al formato de archivo Parquet
- Apache Parquet es un formato de almacenamiento de datos columnar que permite especificar claramente el tipo de datos de cada columna
- Puede almacenar datos en forma de lista (arreglos
float32), por lo que es adecuado para guardar embeddings
- Ofrece mejor rendimiento de guardado y carga que CSV, y permite cargar selectivamente solo parte de los datos
- Incluye compresión, pero los datos de embeddings tienen poca redundancia, así que el efecto de compresión es limitado
Uso de archivos Parquet en Python
- Guardado y carga de archivos Parquet con pandas:
df = pd.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])
df
- pandas no maneja eficientemente datos anidados (listas), y los convierte a
object de numpy
- Al convertirlos a arreglos numpy, se necesita una operación adicional (
np.vstack()), lo que puede degradar el rendimiento
- Guardado y carga de archivos Parquet con polars:
df = pl.read_parquet("mtg-embeddings.parquet", columns=["name", "embedding"])
df
- polars conserva los arreglos
float32 tal cual, y al llamar to_numpy() puede devolver de inmediato un arreglo numpy 2D
- Con la opción
allow_copy=False, se pueden evitar copias innecesarias de datos
embeddings = df["embedding"].to_numpy(allow_copy=False)
- Al agregar nuevos embeddings, también se pueden guardar de forma simple añadiendo una columna
df = df.with_columns(embedding=embeddings)
df.write_parquet("mtg-embeddings.parquet")
Búsqueda de similitud y filtrado con Parquet + polars
- Se puede ejecutar búsqueda de similitud después de filtrar solo los datos que cumplan ciertas condiciones
- Ejemplo: encontrar cartas similares a una carta específica (
query_embed), pero buscando solo cartas del tipo 'Sorcery' y que incluyan el color 'Black'
df_filter = df.filter(
pl.col("type").str.contains("Sorcery"),
pl.col("manaCost").str.contains("B"),
)
embeddings_filter = df_filter["embedding"].to_numpy(allow_copy=False)
idx, _ = fast_dot_product(query_embed, embeddings_filter, k=4)
related_cards = df_filter[idx]
- El tiempo promedio de ejecución es 1.48ms: 37% más lento que buscar en todos los datos, pero sigue siendo rápido
Alternativas para procesar datos vectoriales a gran escala
- El enfoque de Parquet y
dot product es suficiente para manejar cientos de miles de embeddings
- Si se trabaja con datasets más grandes, puede ser necesario usar una base de datos vectorial
- Como alternativa, se puede usar sqlite-vec sobre SQLite para añadir búsqueda vectorial y filtrado
Conclusión
- Una base de datos vectorial no es obligatoria
- La combinación Parquet + polars es una alternativa potente para almacenar, buscar y filtrar embeddings de forma eficiente
- Especialmente en proyectos pequeños, usar archivos Parquet puede ser más rápido y más rentable
- Según el proyecto, es importante elegir la solución adecuada entre Parquet y una base de datos vectorial
- En el repositorio de GitHub se pueden revisar el código y los datos
1 comentarios
Opiniones de Hacker News
El problema con Parquet es que es estático. No es adecuado cuando se necesitan escrituras y actualizaciones continuas. Aun así, obtuve buenos resultados usando archivos Parquet con DuckDB y almacenamiento de objetos. Los tiempos de carga son rápidos
numpy float32como bytes y luego decodificarlos de nuevo como arreglos de numpyfloat32. Toma unos 2 ms para alrededor de 20,000 elementos, lo cual es más rápido que LanceDB. En colecciones más grandes, Lance podría ganar. Pero en mi caso de uso, cada usuario tiene su propio archivo SQLite dedicado, así que funciona bienMuy buen artículo. He disfrutado tu trabajo durante mucho tiempo. Para quienes se lancen a una implementación con SQLite, vale la pena agregar que DuckDB empezó a incorporar algunas funciones de similitud vectorial que leen Parquet y cubren perfectamente este caso de uso
Sigo sin amar los dataframes, pero Polars es muchísimo mejor que pandas
Revisen usearch de Unum. Le gana a cualquier cosa y es muy fácil de usar. Hace exactamente lo que necesitas
Si quieres probarlo, puedes cargarlo de forma diferida desde HF y aplicar filtrado
POLARS_MAX_THREADSa un Ray Actor para ajustarlo según el nivel de saturación del nodo únicoHay muchos hallazgos excelentes
resume.jsona una versión de texto completo y luego genera los embeddings. Los resultados parecen mejores, pero no he visto opiniones concretas sobre estoEn la documentación de Vespa hay un truco elegante que convierte el vector a binario y luego usa una representación hexadecimal
Polars + Parquet es excelente en portabilidad y rendimiento. Esta publicación se enfocó en la portabilidad en Python, pero Polars también tiene una API de Rust fácil de usar que permite incrustar el motor en muchos lugares
Soy muy fan de Polars, pero no había considerado usarlo para almacenar embeddings (había estado experimentando con sqlite-vec). De verdad parece una idea muy interesante
Recomiendo lancedb como otra biblioteca con gran rendimiento y características como indexación de texto completo y versionado de cambios