WordLlama 📝🦙

WordLlama 是一个快速、轻量级的自然语言处理工具包，专为模糊去重、相似度计算、排序、聚类以及语义文本分割等任务而设计。它在推理时依赖极少，并针对 CPU 硬件进行了优化，因此非常适合部署在资源受限的环境中。

新闻与更新 🔥

• 2025-02-01 支持标准库函数（sorted/min/max）的可调用接口
• 2025-01-04 我们很高兴地宣布支持 model2vec 静态嵌入。更多信息请参阅：Model2Vec
• 2024-10-04 添加了语义分割推理算法。请查看我们的技术概述。

目录

• 快速入门
• 功能特性
• 什么是 WordLlama？
• MTEB 测试结果
• 速度有多快？
• 使用方法
  • 文本嵌入
  • 标准库排序/最小值/最大值
  • 计算相似度
  • 文档排序
  • 模糊去重
  • 聚类
  • 过滤
  • Top-K 检索
  • 语义文本分割
  • 加载 Model2Vec
  • 推理类
• 训练笔记
• 路线图
• 提取词元嵌入
• 社区项目
• 引用
• 许可证

快速入门

通过 pip 安装 WordLlama：

pip install wordllama

加载默认的 256 维模型：

from wordllama import WordLlama

# 加载默认的 WordLlama 模型
wl = WordLlama.load()

query = "机器学习方法"
candidates = [
    "神经科学基础",
    "神经网络导论",
    "在家烹饪美味意大利面",
    "哲学导论：逻辑学",
]

# 返回一个 Callable[[str], float] 函数
sim_key = wl.key(query)

# 对候选列表按相似度从高到低排序
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)

# 最相似的候选
best_candidate = max(candidates, key=sim_key)

# 打印结果
print("排序后的候选列表：")
for i, candidate in enumerate(sorted_candidates, 1):
    print(f"{i}. {candidate} (得分: {sim_key(candidate):.4f})")

print(f"\n最佳匹配：{best_candidate} (得分: {sim_key(best_candidate):.4f})")

# 排序后的候选列表：
# 1. 神经网络导论 (得分: 0.3414)
# 2. 神经科学基础 (得分: 0.2115)
# 3. 哲学导论：逻辑学 (得分: 0.1067)
# 4. 在家烹饪美味意大利面 (得分: 0.0045)
#
# 最佳匹配：神经网络导论 (得分: 0.3414)

功能特性

• 快速嵌入：通过简单的词元查找和平均池化，高效生成文本嵌入。
• 相似度计算：计算文本之间的余弦相似度。
• 排序：根据文档与查询的相似度对文档进行排序。
• 模糊去重：基于相似度阈值去除重复文本。
• 聚类：使用 KMeans 聚类算法将文档分组。
• 过滤：根据文档与查询的相似度筛选文档。
• Top-K 检索：检索与查询最相似的前 K 个文档。
• 语义文本分割：将文本分割成语义连贯的块。
• 二进制嵌入：支持二进制嵌入和汉明相似度，以实现更快的计算。
• 套娃式表示：可根据需要截断嵌入维度，提供更大的灵活性。
• 低资源需求：针对 CPU 推理进行了优化，依赖项极少。

什么是 WordLlama？

WordLlama 是一种自然语言处理工具，它复用了大型语言模型（LLMs）中的组件，以创建高效且紧凑的词表示，类似于 GloVe、Word2Vec 或 FastText。

WordLlama 首先从最先进的 LLMs（例如 LLaMA 2、LLaMA 3 70B）中提取词元嵌入码本，然后在一个通用的嵌入框架内训练一个小规模的无上下文模型。这种方法生成了一个轻量级模型，在所有 MTEB 基准测试上都优于传统的词模型（如 GloVe 300d），同时体积也小得多（例如，256 维的默认模型仅 16MB）。

WordLlama 的主要特点包括：

1. 套娃式表示：允许根据需要截断嵌入维度，从而在模型大小和性能之间取得灵活平衡。
2. 低资源需求：采用简单的词元查找结合平均池化，可在 CPU 上快速运行，无需 GPU。
3. 二进制嵌入：使用直通估计器训练的模型可以打包成小型整数数组，以实现更快的汉明距离计算。
4. 纯 NumPy 推理：轻量级的推理流程完全依赖 NumPy，便于部署和集成。

由于其速度快、体积小且便携，WordLlama 可作为探索性分析和实用应用程序的多功能工具，例如用于评估 LLM 输出或在多跳式、代理式工作流中执行预处理任务。

MTEB 评测结果

下表展示了 WordLlama 模型与其他类似模型的性能对比。

指标	WL64	WL128	WL256 (X)	WL512	WL1024	GloVe 300d	Komninos	all-MiniLM-L6-v2
聚类	30.27	32.20	33.25	33.40	33.62	27.73	26.57	42.35
重排序	50.38	51.52	52.03	52.32	52.39	43.29	44.75	58.04
分类	53.14	56.25	58.21	59.13	59.50	57.29	57.65	63.05
对分类	75.80	77.59	78.22	78.50	78.60	70.92	72.94	82.37
STS	66.24	67.53	67.91	68.22	68.27	61.85	62.46	78.90
CQA DupStack	18.76	22.54	24.12	24.59	24.83	15.47	16.79	41.32
SummEval	30.79	29.99	30.99	29.56	29.39	28.87	30.49	30.81

WL64 至 WL1024：WordLlama 模型，嵌入维度从 64 到 1024 不等。

注：l2_supercat 是一个 LLaMA 2 词汇表模型。为了训练该模型，我们在移除额外的特殊标记后，将来自多个模型（包括 LLaMA 2 70B 和 phi 3 medium）的代码本拼接在一起。由于多个模型使用了 LLaMA 2 的分词器，它们的代码本可以拼接并一起训练。最终得到的模型性能与训练 LLaMA 3 70B 代码本相当，但规模仅为前者的四分之一（32k 词汇量 vs. 128k 词汇量）。

其他模型

• 基于 LLaMA 3 的：l3_supercat
• 结果

运行速度如何？ :zap:

ag_news 数据集中的 8,000 条文档

• 单核性能（CPU），i9 第 12 代，DDR4 3200
• NVIDIA A4500（GPU）

使用方法

文本嵌入

加载预训练的嵌入并向文本进行嵌入：

from wordllama import WordLlama

# 加载预训练的嵌入（截断维度至 64）
wl = WordLlama.load(trunc_dim=64)

# 向文本进行嵌入
embeddings = wl.embed(["The quick brown fox jumps over the lazy dog", "And all that jazz"])
print(embeddings.shape)  # 输出：(2, 64)

标准库示例

通过 .key(query) 返回一个可调用函数。

query = "Machine learning methods"
candidates = [
    "Foundations of neural science",
    "Introduction to neural networks",
    "Cooking delicious pasta at home",
    "Introduction to philosophy: logic",
]

# 返回一个 Callable[[str], float] 函数
sim_key = wl.key(query)

# 对候选者按相似度排序，最相似的排在前面
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)

# 最相似的候选者
best_candidate = max(candidates, key=sim_key)

# 打印结果
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)  # 输出：例如 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)
# 输出：
# [
#   ('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)
# 输出：
# ['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)
# 输出：
# [2, 0, 1, 1], 0.4150

过滤

根据文档与查询的相似度进行过滤：

filtered_docs = wl.filter(query, candidates, threshold=0.3)
print(filtered_docs)
# 输出：
# ['I went to the vehicle']

Top-K 检索

检索与查询最相似的前 K 个文档：

top_docs = wl.topk(query, candidates, k=2)
print(top_docs)
# 输出：
# ['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)))
# 输出：[1055, 1055, 1187]

请注意，目标大小也是最大大小。.split() 功能会尝试聚合不超过 target_size 的部分，同时保留文本的顺序以及句子和尽可能多的段落结构。它使用 wordllama 嵌入来找到更自然的分割点。因此，输出中会出现一系列大小不一的块，但都不会超过目标大小。

推荐的目标大小范围是 512 到 2048 个字符，默认大小为 1536。如果需要更大的块，可能需要在分割后进行批量处理，并且通常会由多个语义块组合而成。

更多信息请参阅：技术概述

加载 Model2Vec

wl = WordLlama.list_configs()
# 配置名称字典

wl = WordLlama.load_m2v("potion_base_8m") # 256 维模型
wl = WordLlama.load_m2v("m2v_multilingual") # 多语言模型

Model2Vec 是一种使用 PCA 创建静态嵌入的不同方法。 值得注意的是，他们已经开发出多语言模型和基于 GloVe 的模型，在词语相似度任务中表现优异。

欢迎前往 Hugging Face 查看！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 模型是在单张 A100 GPU 上以 512 的批量大小训练了 12 小时。

路线图

• 更多示例笔记本：
  • DSPy 评估器
  • 检索增强生成（RAG）流水线

开发

本地开发步骤如下：

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。

提取 Token 嵌入

要从模型中提取 Token 嵌入，请确保已同意用户协议并使用 Hugging Face CLI 登录（适用于 LLaMA 模型）。然后可以使用以下代码片段：

from wordllama.extract.extract_safetensors import extract_safetensors

# 提取指定配置的嵌入
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
# （开始训练过程）
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
# （按每个 Matryoshka 维度保存一个模型）

社区项目

• Gradio 演示 HF Space
• CPU-ish RAG

引用

如果您在研究或项目中使用 WordLlama，请考虑按照以下方式引用：

@software{miller2024wordllama,
  author = {Miller, D. Lee},
  title = {WordLlama: 从大型语言模型中回收的 Token 嵌入},
  year = {2024},
  url = {https://github.com/dleemiller/wordllama},
  version = {0.3.9}
}

许可证

本项目采用 MIT 许可证授权。

WordLlama 快速上手指南

WordLlama 是一个轻量级、高速的 NLP 工具包，专为模糊去重、相似度计算、排序、聚类和语义文本分割等任务设计。它仅依赖 CPU 运行，推理时无需重型框架，非常适合资源受限的环境。

环境准备

• 操作系统：Linux, macOS, Windows
• Python 版本：Python 3.8+
• 硬件要求：仅需 CPU（已针对 CPU 推理优化），无需 GPU
• 核心依赖：numpy（安装时会自动处理）

安装步骤

使用 pip 直接安装官方版本：

pip install wordllama

提示：国内开发者若遇到下载速度慢的问题，可使用清华或阿里镜像源加速安装：

pip install wordllama -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

以下示例展示如何加载默认模型并对候选文本进行相似度排序。

1. 加载模型与相似度排序

from wordllama import WordLlama

# 加载默认的 256 维 WordLlama 模型
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",
]

# 生成一个可调用函数，用于计算与 query 的相似度
sim_key = wl.key(query)

# 按相似度降序排序
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)

# 获取最相似的候选项
best_candidate = max(candidates, key=sim_key)

# 打印结果
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)

2. 其他常用功能速览

除了排序，WordLlama 还支持多种开箱即用的功能：

• 计算相似度：wl.similarity("文本 A", "文本 B")
• 文档排序：wl.rank(query, candidates)
• 模糊去重：wl.deduplicate(candidates, threshold=0.5)
• 聚类：wl.cluster(candidates, k=3)
• Top-K 检索：wl.topk(query, candidates, k=2)
• 语义文本分割：wl.split(long_text, target_size=1536)

3. 自定义维度加载

如需更小的模型体积或更快的速度，可截断嵌入维度（例如降至 64 维）：

# 加载并截断至 64 维
wl = WordLlama.load(trunc_dim=64)
embeddings = wl.embed(["The quick brown fox", "And all that jazz"])
print(embeddings.shape)  # 输出：(2, 64)