- WordLlama es un toolkit de NLP rápido y ligero que reutiliza embeddings de tokens de un LLM para hacer deduplicación difusa, cálculo de similitud, ranking, clustering y segmentación semántica de texto
- La inferencia se basa principalmente en búsqueda de tokens y average pooling, y destaca una tubería ligera que puede ejecutarse solo con NumPy junto con optimización para CPU
- El modelo base es de 256 dimensiones y 16 MB; puede reducir dimensiones con representaciones Matryoshka, y los embeddings binarios permiten cálculos más rápidos con similitud de Hamming
- En la tabla de MTEB, WL64~WL1024 muestran puntajes más altos que GloVe 300d y Komninos en varias métricas, aunque en general quedan por debajo de
all-MiniLM-L6-v2
- Se puede usar con
pip install wordllama y WordLlama.load(), y .key(query) devuelve una función invocable que puede pasarse a funciones de la librería estándar como sorted, min y max
Qué hace WordLlama
- WordLlama es un toolkit ligero para tareas utilitarias de NLP como deduplicación difusa, cálculo de similitud, ranking, clustering y segmentación semántica de texto
- Extrae el codebook de embeddings de tokens de LLMs recientes como LLaMA 2 y LLaMA 3 70B para crear representaciones compactas de palabras, similares a GloVe, Word2Vec y FastText
- Durante la inferencia tiene pocas dependencias y está optimizado para hardware CPU, por lo que es adecuado para despliegues en entornos con recursos limitados
- Gracias a su velocidad y tamaño reducido, puede usarse para análisis exploratorio, evaluadores de salidas de LLM y tareas preparatorias en workflows multi-hop o agentic
Instalación y uso básico
- La instalación se hace con pip
pip install wordllama
- El modelo base de 256 dimensiones se carga con
WordLlama.load()
from wordllama import WordLlama
wl = WordLlama.load()
.key(query) devuelve un Callable[[str], float] que permite ordenar cadenas candidatas por similitud con la consulta o elegir la de mayor puntaje
query = "Machine learning methods"
candidates = [
"Foundations of neural science",
"Introduction to neural networks",
"Cooking delicious pasta at home",
"Introduction to philosophy: logic",
]
sim_key = wl.key(query)
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
best_candidate = max(candidates, key=sim_key)
- En el resultado del ejemplo,
"Introduction to neural networks" queda como la mejor candidata con un puntaje de 0.3414
Funciones principales
- Generación de embeddings: crea embeddings de texto rápidamente con simple búsqueda de tokens y average pooling
- Cálculo de similitud: calcula la cosine similarity entre dos textos
- Ranking de documentos: ordena según la similitud entre la consulta y los documentos candidatos
- Deduplicación difusa: elimina textos duplicados según un umbral de similitud
- Clustering: agrupa documentos con KMeans
- Filtrado: deja solo documentos cuya similitud con la consulta supera un umbral
- Búsqueda Top-K: devuelve los K documentos más similares a la consulta
- Segmentación semántica de texto: divide el texto en chunks semánticamente coherentes
- Embeddings binarios: permiten cálculos más rápidos con similitud de Hamming
- Representaciones Matryoshka: recortan la dimensión del embedding según se necesite para ajustar tamaño del modelo y rendimiento
Estructura del modelo y rendimiento
- WordLlama entrena un modelo pequeño sin contexto dentro de un framework general de embeddings
- El modelo base tiene un tamaño de 256 dimensiones y 16 MB
- La tabla MTEB del README compara WL64, WL128, WL256, WL512 y WL1024 con GloVe 300d, Komninos y
all-MiniLM-L6-v2
- WL256 registra Clustering
33.25, Reranking 52.03, Classification 58.21, Pair Classification 78.22, STS 67.91, CQA DupStack 24.12, SummEval 30.99
- GloVe 300d registra en esos mismos apartados
27.73, 43.29, 57.29, 70.92, 61.85, 15.47, 28.87
all-MiniLM-L6-v2 registra Clustering 42.35, Reranking 58.04, Classification 63.05, Pair Classification 82.37, STS 78.90, CQA DupStack 41.32, SummEval 30.81
- l2_supercat es un modelo con vocabulario de LLaMA 2
- Se entrena concatenando los codebooks de varios modelos, como LLaMA 2 70B y phi 3 medium, tras eliminar tokens especiales adicionales
- Puede entrenarse concatenando juntos los codebooks de varios modelos que usan el tokenizer de LLaMA 2
- Muestra un rendimiento similar al entrenamiento con el codebook de LLaMA 3 70B, mientras que el vocabulario es
32k frente a 128k, es decir, 4 veces más pequeño
- También se ofrece l3_supercat, basado en LLaMA 3
- Hay resultados adicionales en Results
Segmentación semántica de texto
.split() divide textos largos en chunks semánticos
long_text = "Your very long text goes here... " * 100
chunks = wl.split(long_text, target_size=1536)
print(list(map(len, chunks)))
# Output: [1055, 1055, 1187]
target_size es el tamaño objetivo y también el tamaño máximo- El proceso de segmentación intenta preservar el orden del texto, la estructura de las oraciones y, cuando es posible, la estructura de los párrafos
- Usa embeddings de WordLlama para encontrar índices de corte más naturales
- El tamaño de los chunks de salida puede variar dentro del límite de
target_size
- El tamaño objetivo recomendado es de 512~2048 caracteres y el valor por defecto es
1536
- Si se necesitan chunks más grandes, se recomienda agrupar en lotes varios chunks semánticos después de la segmentación
- Hay más detalles en el technical overview
Model2Vec e inferencia directa
- En la actualización del 2025-01-04 se agregó soporte para static embeddings de Model2Vec
- Los modelos de Model2Vec pueden cargarse con
WordLlama.load_m2v()
wl = WordLlama.list_configs()
wl = WordLlama.load_m2v("potion_base_8m") # 256-dim model
wl = WordLlama.load_m2v("m2v_multilingual") # multilingual model
- Model2Vec es otro enfoque para crear embeddings estáticos usando PCA
- Del lado de Model2Vec se menciona que crearon un modelo multilingüe y modelos basados en glove, con buenos puntajes en tareas de similitud entre palabras
- Puede revisarse en minishlab en Hugging Face
WordLlamaInference puede usarse ingresando directamente un arreglo de embeddings estáticos con forma (n_vocab, dim) y un tokenizer, en lugar de usar el cargador
from wordllama import WordLlamaInference
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_pretrained(...)
wl = WordLlamaInference(np_embeddings_ar, tokenizer)
Entrenamiento y extracción de embeddings
- Los modelos de embeddings binarios mostraron mejoras más marcadas en dimensiones altas, y para embeddings binarios se recomiendan 512 o 1024 dimensiones
- El modelo L2 Supercat se entrenó durante 12 horas en una sola GPU A100 con batch size
512
- Para extraer embeddings de tokens desde modelos LLaMA, hay que aceptar el acuerdo de usuario e iniciar sesión en Hugging Face CLI
from wordllama.extract.extract_safetensors import extract_safetensors
extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")
- Los embeddings suelen estar en el primer archivo
safetensors, pero no siempre
- Puede haber un manifest
- Puede ser necesario inspeccionarlo manualmente y encontrarlos
- Para el entrenamiento se usan scripts del repositorio, y hay que añadir un configuration file copiando o modificando uno existente
pip install wordllama[train]
python train.py train --config your_new_config
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
- La etapa de guardado almacena un modelo por cada dimensión Matryoshka
Actualizaciones, roadmap y licencia
- En la actualización del 2025-02-01 se agregó soporte para funciones invocables utilizables en funciones de la librería estándar como
sorted, min y max
- En la actualización del 2024-10-04 se agregó el algoritmo de inferencia para semantic splitting
- En el roadmap aparece la incorporación de notebooks de ejemplo para DSPy evaluator y pipelines de Retrieval-Augmented Generation, es decir, RAG
- Como proyectos de la comunidad están Gradio Demo HF Space y CPU-ish RAG
- La licencia del proyecto es MIT License
1 comentarios
Comentarios de Hacker News
Me gusta mucho lo pequeño que es. Ya tiene ventajas incluso frente al modelo más pequeño de SBERT
Dicho eso, técnicamente parece un enfoque bastante antiguo, y entiendo que hay un compromiso con el rendimiento. Aun así, me pregunto si podría ofrecer cambio de tipo de similitud para cosas como similitud semántica, inferencia en lenguaje natural (NLI) o abstracción de sustantivos
Por ejemplo, al agrupar artículos de periódico en categorías como “eventos ambientales extremos”, querría que “Freezing” y “Burning” salieran como muy parecidos. Ese sería el caso de MTEB/Sentence-Similarity o del clásico Word2Vec/GloVe. Pero si fueran artículos de química, deberían salir casi como opuestos, y a veces también quisiera ver la relación causal entre dos cosas usando embeddings de inferencia en lenguaje natural
Los dos últimos tipos de embedding son enfoques relativamente recientes, posteriores a 2019, así que creo que ahí hay más oportunidad técnica. La familia antigua de MTEB/similitud semántica ya era suficiente para muchos usos desde 2014, y mejoró mucho en 2019 con mini-lm-v2 y similares
Los tres tipos de embedding de arriba también son posibles con SBERT, pero como la dimensionalidad es grande y los modelos también, cargar varios modelos por tipo implica bastante costo de recursos. Los modelos generativos de embeddings, E5 y los de inferencia en lenguaje natural son grandes y a menudo requieren como 6GB
Me gustaría ver cómo rinde al entrenarse con un solo tipo de similitud. No estoy seguro de que haya otra forma de manejar esto sin calcular contexto. Tal vez haya que cambiar de modelo, pero eso por sí mismo no sería un gran problema
No lo digo para menospreciarlo; es importante entender este trabajo en contexto. Aquí el contexto es que, al profundizar en cómo funcionan los LLM, uno descubre que también tienen embeddings, y desde esa perspectiva resulta más natural avanzar un paso jugando con esos embeddings que volver a revisar todo el estado del arte del campo de embeddings
Pensar que palabras con significado opuesto deberían tener embeddings opuestos es un malentendido muy común. En la práctica, las palabras antónimas también comparten muchas cosas. “burning” y “freezing” están relacionadas con temperatura y física, ambas son palabras en inglés, pueden usarse como verbo, sustantivo y adjetivo, y además están bien escritas. Todo ese tipo de rasgos también entra en el embedding
Los embeddings pueden contener mucha información semántica según los datos de entrenamiento y la función objetivo, y pueden usarse de forma independiente para varias tareas útiles
Hace tiempo usé embeddings del codificador de texto de CLIP para reforzar prompts y hacer que coincidieran mejor con la imagen correspondiente. Por ejemplo, si el prompt contenía “building”, buscaba vecinos más cercanos como “concrete” o “underground” en la matriz de embeddings, y reemplazaba o añadía esas palabras después del término original. En experimentos limitados, eso mejoró la recuperación en la mayoría de las consultas
https://www.marqo.ai/blog/generalized-contrastive-learning-f...
Mirando el tamaño de los embeddings de tokens en wordllama, también podría ayudar a identificar qué tokens importantes conviene reforzar. Aun así, si se entrena con datos seleccionados específicamente para esta tarea, probablemente funcionaría mucho mejor
Me pregunto si hay planes para idiomas aparte del inglés. Para el francés parece que sería una herramienta perfecta
Ya he probado un poco de entrenamiento con modelos de la familia Mistral, así que probablemente empezaría por ahí para un corpus en francés
Si abres un issue, puedo intentarlo cuando tenga tiempo
Para corpus grandes, por ejemplo cuando cada oración se trata como un documento en un conjunto de más de 10,000 oraciones, se podrían obtener resultados parecidos agrupando vectores dispersos TF-IDF con k-means
Aunque parece que esta herramienta sí tiene bastantes utilidades para acelerar más la parte de k-means con cosas como binarización. Planeo hacer benchmarks durante las próximas semanas
Hace algunos años hice una colección de juegos de lenguaje que usan funciones parecidas: https://github.com/Hellisotherpeople/Language-games
https://github.com/plasticityai/magnitude
Me pregunto si alguien ha pensado en usar embeddings para resolver Little Alchemy. #sample-use
Se ve genial. Me pregunto si tiene alguna ventaja frente a los modelos mini-lm. En la mayoría de las tareas de MTEB parece que mini-lm sale mejor, así que me pregunto si habrá alguna mejora en cosas como velocidad de inferencia
Tampoco estaba apuntando al mejor rendimiento posible del momento. Es un modelo con restricciones bastante fuertes para reducir dependencias, tamaño y requisitos de hardware, y para ganar velocidad
Incluso como modelo de embeddings de palabras, es bastante liviano. Normalmente esos modelos tienen vocabularios mucho más grandes y pesan varios GB
Parece muy útil para desarrollo de juegos
Muestra muy bien cuánta carga semántica hay en los propios tokens
¿Se podría hacer como una extensión de PostgreSQL?