8.0
热点
AI SCORE
开源项目2024-09-15 11:25
Wordllama：LLM Token Embeddings 实用工具集
Hacker News · github.com#向量化#Embeddings
Editor brief · 编辑速览
开源项目展示 token embeddings 在向量化处理中的多种应用场景，提供即插即用的代码示例。
WordLlama 是一个快速、轻量级的自然语言处理（NLP）工具包，专为模糊去重、相似度计算、排序、聚类和语义文本分割等任务设计。它具有最少的推理时依赖，并针对 CPU 硬件进行了优化，适合在资源受限的环境中部署。
2025-02-01 为标准库函数（sorted/min/max）添加可调用接口
2025-01-04 我们很高兴地宣布支持 model2vec 静态嵌入。参见：Model2Vec
2024-10-04 添加了语义分割推理算法。查看我们的技术概览。
通过 pip 安装 WordLlama：
pip install wordllama
加载默认的 256 维模型：
from wordllama import WordLlama
# Load the default WordLlama model
wl = WordLlama.load()
query = "Machine learning methods"
candidates = [
"Foundations of neural science",
"Introduction to neural networks",
"Cooking delicious pasta at home",
"Introduction to philosophy: logic",
]
# Returns a Callable[[str], float] function
sim_key = wl.key(query)
# Sort candidates, most similar first
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
# Most similar candidate
best_candidate = max(candidates, key=sim_key)
# Print the results
print("Ranked Candidates:")
for i, candidate in enumerate(sorted_candidates, 1):
print(f"{i}. {candidate} (Score: {sim_key(candidate):.4f})")
print(f"\nBest Match: {best_candidate} (Score: {sim_key(best_candidate):.4f})")
# Ranked Candidates:
# 1. Introduction to neural networks (Score: 0.3414)
# 2. Foundations of neural science (Score: 0.2115)
# 3. Introduction to philosophy: logic (Score: 0.1067)
# 4. Cooking delicious pasta at home (Score: 0.0045)
#
# Best Match: Introduction to neural networks (Score: 0.3414)
快速嵌入：使用简单的令牌查找和平均池化，高效地生成文本嵌入。
相似度计算：计算文本之间的余弦相似度。
排序：根据与查询的相似度对文档进行排序。
模糊去重：基于相似度阈值移除重复文本。
聚类：使用 KMeans 聚类将文档分组。
筛选：根据与查询的相似度对文档进行筛选。
Top-K 检索：检索与查询最相似的 Top-K 个文档。
语义文本分割：将文本分割成语义一致的块。
二进制嵌入：支持二进制嵌入和汉明相似度，实现更快的计算。
俄罗斯套娃表征：根据需要截断嵌入维度以获得灵活性。
低资源需求：针对 CPU 推理优化，依赖最少。
WordLlama 是一个自然语言处理（NLP）实用工具，它重新利用大语言模型（LLM）的组件来创建高效而紧凑的词表征，类似于 GloVe、Word2Vec 或 FastText。
WordLlama 通过从最先进的 LLM（例如 LLaMA 2、LLaMA 3 70B）中提取令牌嵌入代码簿，在通用嵌入框架内训练一个小型无上下文模型。这种方法生成的轻量级模型在所有 MTEB 基准测试中都优于传统词模型（如 GloVe 300d），同时体积明显更小（例如 256 维的默认模型仅 16MB）。
WordLlama 的主要特性包括：
俄罗斯套娃表征：允许根据需要截断嵌入维度，在模型大小和性能之间提供灵活性。
低资源需求：使用简单的令牌查找和平均池化，能在 CPU 上快速运行，无需 GPU。
二进制嵌入：使用直通估计器训练的模型可以打包成小整数数组，用于更快的汉明距离计算。
仅 NumPy 推理：轻量级推理管道仅依赖 NumPy，便于部署和集成。
由于其快速和便携的体积，WordLlama 是一个用途广泛的工具，适用于探索性分析和实用应用，如 LLM 输出评估器或多跳或代理工作流中的预备任务。
以下表格展示了 WordLlama 模型与其他类似模型的性能对比。
WL64 至 WL1024：WordLlama 模型，嵌入维度范围从 64 到 1024。
注：l2_supercat 是一个 LLaMA 2 词汇表模型。为了训练该模型，我们在删除额外特殊令牌后，连接了来自多个模型的代码簿，包括 LLaMA 2 70B 和 phi 3 medium。由于多个模型使用了 LLaMA 2 分词器，它们的代码簿可以连接并一起训练。所得模型的性能与训练 LLaMA 3 70B 代码簿相当，同时体积缩小了 4 倍（32k vs 128k 词汇表）。
基于 LLaMA 3 的：l3_supercat
来自 ag_news 数据集的 8000 个文档
单核性能（CPU），i9 12 代，DDR4 3200
from wordllama import WordLlama
# Load pre-trained embeddings (truncate dimension to 64)
wl = WordLlama.load(trunc_dim=64)
# Embed text
embeddings = wl.embed(["The quick brown fox jumps over the lazy dog", "And all that jazz"])
print(embeddings.shape) # Output: (2, 64)
query = "Machine learning methods"
candidates = [
"Foundations of neural science",
"Introduction to neural networks",
"Cooking delicious pasta at home",
"Introduction to philosophy: logic",
]
# Returns a Callable[[str], float] function
sim_key = wl.key(query)
# Sort candidates, most similar first
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
# Most similar candidate
best_candidate = max(candidates, key=sim_key)
# Print the results
print("Ranked Candidates:")
for i, candidate in enumerate(sorted_candidates, 1):
print(f"{i}. {candidate} (Score: {sim_key(candidate):.4f})")
print(f"\nBest Match: {best_candidate} (Score: {sim_key(best_candidate):.4f})")
# Ranked Candidates:
# 1. Introduction to neural networks (Score: 0.3414)
# 2. Foundations of neural science (Score: 0.2115)
# 3. Introduction to philosophy: logic (Score: 0.1067)
# 4. Cooking delicious pasta at home (Score: 0.0045)
#
# Best Match: Introduction to neural networks (Score: 0.3414)
计算两个文本之间的相似度：
similarity_score = wl.similarity("I went to the car", "I went to the pawn shop")
print(similarity_score) # Output: e.g., 0.0664
根据文档与查询的相似度对其进行排序：
query = "I went to the car"
candidates = ["I went to the park", "I went to the shop", "I went to the truck", "I went to the vehicle"]
ranked_docs = wl.rank(query, candidates, sort=True, batch_size=64)
print(ranked_docs)
# Output:
# [
# ('I went to the vehicle', 0.7441),
# ('I went to the truck', 0.2832),
# ('I went to the shop', 0.1973),
# ('I went to the park', 0.1510)
# ]
基于相似度阈值移除重复文本：
deduplicated_docs = wl.deduplicate(candidates, return_indices=False, threshold=0.5)
print(deduplicated_docs)
# Output:
# ['I went to the park',
# 'I went to the shop',
# 'I went to the truck']
使用 KMeans 聚类将文档分组：
labels, inertia = wl.cluster(candidates, k=3, max_iterations=100, tolerance=1e-4, n_init=3)
print(labels, inertia)
# Output:
# [2, 0, 1, 1], 0.4150
根据与查询的相似度筛选文档：
filtered_docs = wl.filter(query, candidates, threshold=0.3)
print(filtered_docs)
# Output:
# ['I went to the vehicle']
检索与查询最相似的 Top-K 个文档：
top_docs = wl.topk(query, candidates, k=2)
print(top_docs)
# Output:
# ['I went to the vehicle', 'I went to the truck']
将文本分割成语义块：
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]
注意，目标大小也是最大大小。.split() 功能尝试将部分聚合到 target_size，但会保留文本的顺序以及句子和尽可能多的段落结构。它使用 wordllama 嵌入来定位更自然的分割索引。因此，输出中的块大小范围会在目标大小以内。
建议的目标大小为 512 到 2048 个字符，默认大小为 1536。需要大得多的块应该在分割后进行批处理，通常会从已经分割的多个语义块中进行聚合。
更多信息请参见：technical overview
wl = WordLlama.list_configs()
# dict of config names
wl = WordLlama.load_m2v("potion_base_8m") # 256-dim model
wl = WordLlama.load_m2v("m2v_multilingual") # multilingual model
Model2Vec 是一种使用 PCA 创建静态嵌入的不同方法。值得注意的是，他们已经生成了多语言模型和基于 glove 的模型，在词相似度任务中得分很好。
在 huggingface 上查看它们！minishlab
from wordllama import WordLlamaInference
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_pretrained(...)
wl = WordLlamaInference(np_embeddings_ar, tokenizer)
推理类可以直接与自备的静态嵌入数组（n_vocab, dim）一起使用，而不是使用加载器。
二进制嵌入模型在较高维度时显示更明显的改进，二进制嵌入建议使用 512 或 1024 维。
L2 Supercat 模型使用批量大小 512 在单个 A100 GPU 上训练 12 小时。
git clone https://github.com/dleemiller/WordLlama.git
cd WordLlama
pip install uv
uv sync --all-extras
uv run python setup.py build_ext --inplace
uv run pytest
查看 Makefile 了解常见的开发命令。
要从模型中提取令牌嵌入，请确保您已同意用户协议并使用 Hugging Face CLI 登录（对于 LLaMA 模型）。然后您可以使用以下代码片段：
from wordllama.extract.extract_safetensors import extract_safetensors
# Extract embeddings for the specified configuration
extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")
提示：嵌入通常在第一个 safetensors 文件中，但并非总是如此。有时有清单；有时你必须检查并找出来。
对于训练，使用 GitHub 存储库中的脚本。您必须添加一个配置文件（将现有配置复制/修改到文件夹中）。
pip install wordllama[train]
python train.py train --config your_new_config
# (Training process begins)
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
# (Saves one model per Matryoshka dimension)
如果您在研究或项目中使用 WordLlama，请考虑按如下方式引用它：
@software{miller2024wordllama,
author = {Miller, D. Lee},
title = {WordLlama: Recycled Token Embeddings from Large Language Models},
year = {2024},
url = {https://github.com/dleemiller/wordllama},
version = {0.3.9}
}
该项目在 MIT 许可证下授权。