文章目录
论文: LUXICAL: HIGH-SPEED LEXICAL-DENSE TEXT EMBEDDINGS
luxical是datology公司发布的最新的高效的文本embedding模型, 用于将一篇文档压缩成一个1维的vector, 方便文档聚类. 这在大模型的数据预处理中很有用.
模型: huggingface, 模型的key为:
python
from safetensors.torch import load_file as load_safetensors
state_dict = load_safetensors(weight_path)
print(state_dict.keys())
['embedder.ngram_keys',
'embedder.ngram_vals',
'embedder.num_layers',
'embedder.recognized_ngrams',
'embedder.version',
'embedder.idf_values',
'embedder.nn_layer_0',
'embedder.nn_layer_1',
'embedder.nn_layer_2',
'embedder.nn_layer_3',
'embedder.tokenizer'
]编码的具体流程为:
- 先用tokenizer处理成input_ids
- 然后计算这些input_ids的bow (bag of words)
- 通过tf-idf调整权重
bow:
python
it = pool.imap(
lambda tokens: sparse_count_ngram_in_document(
max_ngram_length=max_ngram_length,
tokens=tokens,
ngram_hash_to_idx=ngram_hash_to_idx,
),
tokenized_docs,
chunksize=worker_chunksize,
)串行逻辑:
python
it = (
sparse_count_ngram_in_document(
max_ngram_length=max_ngram_length,
tokens=tokens,
ngram_hash_to_idx=ngram_hash_to_idx,
)
for tokens in tokenized_docs
)具体执行步骤:
python
def sparse_count_ngram_in_document(
max_ngram_length: int,
tokens: NDArray[T_Token],
ngram_hash_to_idx: NumbaTypedDict[np.int64, np.uint32],
) -> tuple[NDArray[np.uint32], NDArray[np.uint32]]:
"""Gets the indices and counts of the ngrams in the document."""
count_dict = NumbaTypedDict.empty(key_type=numba.uint32, value_type=numba.uint32)
empty_token_val = np.iinfo(tokens.dtype).max
assert not np.any(tokens == empty_token_val), "reserved value"
ng = np.full(max_ngram_length, empty_token_val, dtype=tokens.dtype)
for ngram_length in range(1, max_ngram_length + 1):
num_ngrams_of_this_len = len(tokens) - ngram_length + 1
for i in range(num_ngrams_of_this_len):
ng[:] = empty_token_val
ng[:ngram_length] = tokens[i : i + ngram_length]
ngh = fnv1a_hash_array_to_int64(ng)
if ngh in ngram_hash_to_idx:
idx = ngram_hash_to_idx[ngh]
if idx in count_dict:
count_dict[idx] = np.uint32(count_dict[idx] + np.uint32(1))
else:
count_dict[idx] = np.uint32(1)
idx_array = np.empty(len(count_dict), dtype=np.uint32)
count_array = np.empty(len(count_dict), dtype=np.uint32)
for i, (k, v) in enumerate(count_dict.items()):
idx_array[i] = k
count_array[i] = v
return idx_array, count_array代码按下面顺序执行。
1. 初始化计数字典
count_dict = NumbaTypedDict.empty(
key_type=numba.uint32,
value_type=numba.uint32,
)
初始状态：
count_dict = {}
它保存：
词表特征 idx -> 在当前文档中的出现次数
2. 设置 padding 值
empty_token_val = np.iinfo(tokens.dtype).max
因为输入是 uint32：
empty_token_val = 4294967295
然后检查输入没有使用这个保留值：
assert not np.any(tokens == empty_token_val)
3. 创建固定长度工作数组
ng = np.full(5, empty_token_val, dtype=np.uint32)
此时：
ng = [E, E, E, E, E]
4. 第一次外层循环
ngram_length = 1
num_ngrams_of_this_len = 7 - 1 + 1 # 7
内层第一次：
i = 0
ng[:] = E
ng[:1] = tokens[0:1]
得到：
ng = [1045, E, E, E, E]
接着：
ngh = fnv1a_hash_array_to_int64(ng)
如果该 hash 在词表中：
idx = ngram_hash_to_idx[ngh]
第一次命中时：
count_dict[idx] = 1
内层第二次：
i = 1
ng[:] = E
ng[:1] = tokens[1:2]
得到：
ng = [2066, E, E, E, E]
依次处理完 7 个 unigram。
5. 第二次外层循环
ngram_length = 2
num_ngrams_of_this_len = 7 - 2 + 1 # 6
第一次：
i = 0
ng[:] = E
ng[:2] = tokens[0:2]
得到：
ng = [1045, 2066, E, E, E]
第二次：
i = 1
ng[:] = E
ng[:2] = tokens[1:3]
得到：
ng = [2066, 5983, E, E, E]
依次处理 6 个 bigram。
后续同理：
ngram_length=3 → 5 次
ngram_length=4 → 4 次
ngram_length=5 → 3 次
每次开头都执行：
ng[:] = empty_token_val
这是为了清除上一次循环残留的 token。
6. 累加命中特征
如果多个 ngram 映射到同一个 idx：
if idx in count_dict:
count_dict[idx] += 1
else:
count_dict[idx] = 1
假设最终得到：
count_dict = {
123: 2,
456: 1,
}
7. 转换成数组
idx_array = [123, 456]
count_array = [2, 1]
最后返回：
return idx_array, count_array随后我们要返回ngram特征. 正常情况下对n篇文本, 设词表大小为10, 我们应该返回:
文档0: [0, 0, 1, 0, 0, 2, 0, 0, 0, 0]
文档1: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
文档2: [0, 1, 0, 0, 0, 1, 0, 0, 0, 3]
it 中每个元素对应一篇文档：
(indices, counts)
例如三篇文档返回：
# 文档 0
indices = [2, 5]
counts = [1, 2]
# 文档 1：没有命中词表
indices = []
counts = []
# 文档 2
indices = [1, 5, 9]
counts = [1, 1, 3]
假设循环前：
j_indices_chunks = []
values_chunks = []
indptr = [0]
offset = 0
文档 0
j_indices_chunks.append([2, 5])
values_chunks.append([1, 2])
offset += 2 # offset = 2
indptr.append(2) # [0, 2]
含义是 CSR 第 0 行的数据位于 [0:2]。
文档 1
j_indices_chunks.append([])
values_chunks.append([])
offset += 0 # offset = 2
indptr.append(2) # [0, 2, 2]
indptr 连续两个值相同，表示这一行为空。
文档 2
j_indices_chunks.append([1, 5, 9])
values_chunks.append([1, 1, 3])
offset += 3 # offset = 5
indptr.append(5) # [0, 2, 2, 5]
循环结束后合并：
j_indices = np.concatenate(j_indices_chunks)
# [2, 5, 1, 5, 9]
values = np.concatenate(values_chunks)
# [1, 2, 1, 1, 3]
indptr
# [0, 2, 2, 5]
这三个数组是 CSR 的核心：
values：非零值，即 ngram 出现次数。
j_indices：每个值所在的列，即 ngram 词表下标。
indptr：每一行在前两个数组中的起止位置。
假设词表大小为 10，最终矩阵等价于：
文档0: [0, 0, 1, 0, 0, 2, 0, 0, 0, 0]
文档1: [0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
文档2: [0, 1, 0, 0, 0, 1, 0, 0, 0, 3]
但 CSR 只存储非零项，避免创建巨大的稠密矩阵。因此我们把
[ [0, 0, 1, 0, 0, 2, 0, 0, 0, 0]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
[0, 1, 0, 0, 0, 1, 0, 0, 0, 3]]
弄成稀疏存储格式.
shape = (len(indptr) - 1, vocab_size)
其中：
行数 = 文档数量
列数 = ngram 词表大小
python
for indices, counts in it:
# Update the CSR matrix state.
j_indices_chunks.append(indices)
values_chunks.append(counts)
offset += len(indices)
indptr.append(offset)
pbar.update()
pbar_tok.update(counts.sum())
# Construct the CSR matrix, just as above in the unigram case.
j_indices = np.concatenate(j_indices_chunks)
values = np.concatenate(values_chunks)
vocab_size = len(ngram_hash_to_idx)
matrix = sp.csr_matrix(
(values, j_indices, indptr), shape=(len(indptr) - 1, vocab_size)
)
return matrix接下来我们需要用TF-IDF算法对上面得到的特征进行微调, 这里使用的是原始出现次数作为 TF:
这一步把前面的 BoW 计数矩阵转换成 TF-IDF，并对每篇文档做 L2 归一化。
假设 CSR 矩阵包含：
bow.data # 每个 ngram 的出现次数
bow.indices # 对应的 ngram 词表下标
bow.indptr # 每篇文档在 data/indices 中的范围
例如：
bow.data = [1, 2, 1, 1, 3]
bow.indices = [2, 5, 1, 5, 9]
bow.indptr = [0, 2, 2, 5]
表示：
文档 0: feature 2 出现 1 次，feature 5 出现 2 次
文档 1: 空
文档 2: feature 1 出现 1 次，feature 5 出现 1 次，feature 9 出现 3 次
外层函数
tfidf_data = np.empty_like(bow.data, dtype=np.float32)
创建一个与 bow.data 长度相同的数组：
bow.data: uint32，存放计数
tfidf_data: float32，存放计算后的 TF-IDF
这里只创建新的非零值数组，不创建完整的稠密矩阵。
然后调用：
_fast_tfidf_from_bow(
tfidf_data=tfidf_data,
bow_data=bow.data,
bow_indices=bow.indices,
bow_indptr=bow.indptr,
idf_values=self.idf_values,
)
cast(...) 只用于类型检查，不会在运行时转换或复制 bow.data。
按行处理
num_rows = bow_indptr.shape[0] - 1
因为：
bow_indptr = [0, 2, 2, 5]
所以：
num_rows = 4 - 1 = 3
即三篇文档。
for i in numba.prange(num_rows):
每个 i 对应一篇文档。如果函数使用了类似：
@numba.njit(parallel=True)
那么 prange 可以让不同文档在多个 CPU 核并行处理；否则基本等价于普通 range。
取得当前文档的数据
以文档 0 为例：
start = bow_indptr[0] # 0
end = bow_indptr[1] # 2
row_bow = bow_data[0:2] # [1, 2]
row_indices = bow_indices[0:2] # [2, 5]
row_tfidf = tfidf_data[0:2]
row_tfidf 是 tfidf_data 的视图，对它赋值会直接修改原数组。
应用 IDF
假设：
idf_values[2] = 0.5
idf_values[5] = 1.5
执行：
weighted_val = row_bow[j] * idf_values[row_indices[j]]
得到：
feature 2: TF=1, IDF=0.5 → 1 × 0.5 = 0.5
feature 5: TF=2, IDF=1.5 → 2 × 1.5 = 3.0
所以暂时：
row_tfidf = [0.5, 3.0]
这里使用的是原始出现次数作为 TF，没有额外进行 log(1 + count) 变换。
数学公式是：
TF-IDF(d, feature) = count(d, feature) × idf(feature)
计算 L2 范数
在计算 TF-IDF 的同时累计平方和：
l2_norm_sq += weighted_val * weighted_val
示例：
l2_norm_sq = 0.5² + 3.0²
= 0.25 + 9
= 9.25
然后：
l2_norm = sqrt(9.25) ≈ 3.04138
L2 归一化
row_tfidf /= l2_norm
得到：
0.5 / 3.04138 ≈ 0.1644
3.0 / 3.04138 ≈ 0.9864
归一化后的向量：
row_tfidf = [0.1644, 0.9864]
它满足：
sqrt(0.1644² + 0.9864²) ≈ 1
这样不同长度文档的向量尺度更可比。
如果当前文档为空，或者所有 IDF 权重都为零：
l2_norm_sq == 0
则跳过归一化，避免除以零。
构造最终 CSR 矩阵
return csr_matrix(
(tfidf_data, bow.indices, bow.indptr),
shape=bow.shape,
)
新矩阵保留：
相同的行列形状
相同的非零位置
相同的 indices
相同的 indptr
只把原来的整数计数：
[1, 2, 1, 1, 3]
替换为归一化后的浮点 TF-IDF 值。
整个过程只处理实际存在的非零 ngram，复杂度约为：
O(nnz)
其中 nnz 是所有文档命中的 ngram 特征数量，而不是：
文档数 × 词表大小
因此即使词表有约两百万维，只要每篇文档命中的 ngram 较少，计算和内存开销仍然可控。为什么要用IDF微调?
这里为什么要×idf?
乘以 IDF 是为了降低"到处都出现的常见 ngram"的权重，提高"较少见但更能区分文本的 ngram"的权重。
例如两个 ngram：
"the" 几乎所有文档都有
"machine learning" 只在少量相关文档中出现
如果只用出现次数 TF：
"the" 出现 5 次
"machine learning" 出现 1 次
那么 "the" 权重更高，但它对判断文本主题几乎没帮助。
IDF 通常类似：
IDF(t) = log(文档总数 / 包含特征 t 的文档数)
因此：
常见 ngram → IDF 小
稀有 ngram → IDF 大
乘积：
TF-IDF = 出现次数 × IDF
假设：
"the": TF=5, IDF=0.1 → TF-IDF=0.5
"machine learning": TF=1, IDF=3.0 → TF-IDF=3.0
后者虽然只出现一次，但更能代表文本内容。
在 Luxical 中，idf_values 已经在训练/语料统计阶段算好，并按 ngram_hash_to_idx 的词表下标排列：
idf_values[row_indices[j]]
就是当前 ngram 对应的 IDF。之后再做 L2 归一化，让不同长度文档的向量尺度一致。idf_values 已经预先计算好并保存在模型 checkpoint 中。
加载流程：
python
idf_values = (
state_dict["embedder.idf_values"]
.cpu()
.numpy()
.astype(np.float32, copy=False)
)随后传入：
Embedder(
...
idf_values=idf_values,
)推理时不会重新统计语料，只根据词表下标读取：idf_values[row_indices[j]]
对应关系是：
ngram hash
→ ngram_hash_to_idx 得到词表下标
→ idf_values[词表下标]这些 IDF 值来自模型构建阶段使用的语料统计。推理阶段只是查表并与当前文档中的 ngram 次数相乘。
计算方法为:
I D F ( t ) = l o g ( 文档总数 / 包含特征 t 的文档数 ) IDF(t) = log(文档总数 / 包含特征 t 的文档数) IDF(t)=log(文档总数/包含特征t的文档数)
因此一个token在总语料库中出现的越频繁, 包含特征 t 的文档数越大, 对应的权重越小.
接下来的步骤是将TF-IDF特征转换为dense feature.
python
def csr_matvecs_tiled_unrolled_8(
A_indptr: NDArray[np.int64],
A_indices: NDArray[np.int64],
A_data: NDArray[np.float32],
B_dense: NDArray[np.float32],
Y_out_dense: NDArray[np.float32],
) -> None:
"""Fast multithreaded sparse-by-dense matrix multiplication.
Computes Y_out_dense += A @ B_dense
Uses tiling and manual loop unrolling with block_size 8 for high performance.
"""
BLOCK_SIZE = 8
n_row = A_indptr.shape[0] - 1
n_vecs = B_dense.shape[1]
if n_row == 0:
return
# Process all vectors in blocks. The last block might be smaller.
for k_start in range(0, n_vecs, BLOCK_SIZE):
k_end = min(k_start + BLOCK_SIZE, n_vecs)
# Use the fast, unrolled path if the block is full-sized.
if (k_end - k_start) == BLOCK_SIZE:
for i in numba.prange(n_row):
for jj in range(A_indptr[i], A_indptr[i + 1]):
j = A_indices[jj]
a = A_data[jj]
# Add directly to Y to maintain numerical stability.
# Unrolled for a block of 8.
Y_out_dense[i, k_start + 0] += a * B_dense[j, k_start + 0]
Y_out_dense[i, k_start + 1] += a * B_dense[j, k_start + 1]
Y_out_dense[i, k_start + 2] += a * B_dense[j, k_start + 2]
Y_out_dense[i, k_start + 3] += a * B_dense[j, k_start + 3]
Y_out_dense[i, k_start + 4] += a * B_dense[j, k_start + 4]
Y_out_dense[i, k_start + 5] += a * B_dense[j, k_start + 5]
Y_out_dense[i, k_start + 6] += a * B_dense[j, k_start + 6]
Y_out_dense[i, k_start + 7] += a * B_dense[j, k_start + 7]
# Use a generic tiled loop for the smaller remainder block.
else:
for i in numba.prange(n_row):
for jj in range(A_indptr[i], A_indptr[i + 1]):
j = A_indices[jj]
a = A_data[jj]
for k in range(k_start, k_end):
Y_out_dense[i, k] += a * B_dense[j, k]
@dataclass(frozen=True)
class SparseToDenseEmbedder:
layers: list[NDArray[np.float32]]
@classmethod
def create(
cls,
dims: Sequence[int],
dtype: type[np.floating] = np.float32,
seed: int = 0,
) -> SparseToDenseEmbedder:
assert len(dims) > 1
logger.info(
f"Initializing sparse to dense embedder with dims {dims} and dtype {dtype}."
)
rng = np.random.default_rng(seed)
layers: list[NDArray[np.floating]] = []
for i in range(len(dims) - 1):
kaiming_uniform_scale = np.sqrt(6 / dims[i])
layer = rng.uniform(
-kaiming_uniform_scale,
kaiming_uniform_scale,
size=(dims[i + 1], dims[i]),
).astype(dtype)
layers.append(layer)
return cls(layers)
@property
def output_dim(self) -> int:
return self.layers[-1].shape[0]
@property
def input_dim(self) -> int:
return self.layers[0].shape[1]
@property
def dims(self) -> Sequence[int]:
return [layer.shape[1] for layer in self.layers]
def to_torch(
self,
device: torch.device = torch.device("cpu"),
dtype: torch.dtype = torch.float32,
) -> SparseToDenseEmbedderTorchModule:
return SparseToDenseEmbedderTorchModule(
layers=[torch.tensor(x, device=device, dtype=dtype) for x in self.layers],
)
@classmethod
def from_torch(
cls, module: SparseToDenseEmbedderTorchModule
) -> SparseToDenseEmbedder:
return cls(
layers=[layer.detach().float().cpu().numpy() for layer in module.layers],
)
def __call__(
self,
bow_csr_matrix: csr_matrix,
batch_size: int = 4096,
progress_bar: bool = False,
out: NDArray[np.float32] | None = None,
) -> NDArray[np.float32]:
"""Embed BoW-encoded text into a dense vector space.
Internally batches data to avoid memory overhead. Optionally provides a
`tqdm` progress bar.
"""
assert bow_csr_matrix.ndim == 2
assert bow_csr_matrix.shape[1] == self.layers[0].shape[1] # type: ignore[index]
assert len(self.layers) > 0
num_rows = bow_csr_matrix.shape[0] # type: ignore[index]
# Allocate the final output and chunks for intermediate activations.
chunk_activations = [
np.zeros((batch_size, layer.shape[0]), dtype=np.float32)
for layer in self.layers[:-1]
]
if out is not None:
assert out.shape == (num_rows, self.output_dim)
assert out.dtype == np.float32
else:
out = np.empty((num_rows, self.output_dim), dtype=np.float32)
# Go batch by batch.
with (
tqdm(
total=num_rows,
disable=not progress_bar,
desc="BoW to dense",
unit="row",
unit_scale=True,
) as pbar,
np.testing.suppress_warnings() as sup,
):
# Suppress matmul warnings, which seem to be hair-trigger on some systems.
sup.filter(RuntimeWarning, message=r".*encountered in matmul.*")
for start in range(0, num_rows, batch_size):
# Take the next batch of rows as a floating point matrix.
end = min(start + batch_size, num_rows)
batch_size = end - start
bow_batch = bow_csr_matrix[start:end].astype(np.float32)
# Sparse-by-dense matrix multiplication of the first layer.
chunk_activations[0][:] = 0
csr_matvecs_tiled_unrolled_8(
bow_batch.indptr,
bow_batch.indices,
bow_batch.data,
self.layers[0].T,
chunk_activations[0][:batch_size],
)
# ReLU and normalize.
relu_inplace(chunk_activations[0][:batch_size])
normalize_inplace(chunk_activations[0][:batch_size])
# Do the hidden layers.
for i in range(1, len(self.layers) - 1):
chunk_activations[i][:] = 0
np.matmul(
chunk_activations[i - 1][:batch_size],
self.layers[i].T,
out=chunk_activations[i][:batch_size],
)
relu_inplace(chunk_activations[i][:batch_size])
normalize_inplace(chunk_activations[i][:batch_size])
# Do the final layer.
np.matmul(
chunk_activations[-1][:batch_size],
self.layers[-1].T,
out=out[start:end],
)
normalize_inplace(out[start:end])
pbar.update(batch_size)
return out核心是：用预训练的全连接矩阵，将高维稀疏 TF-IDF 向量投影到低维稠密空间。
假设输入：
TF-IDF CSR: (文档数, ngram 词表大小)
Luxical-One 的词表约两百万维，但每篇文档只有少量非零值。
每个 layer 的形状是：
(output_dim, input_dim)
对行向量计算时使用转置：
output = input @ layer.T
整体流程可以表示为：
稀疏 TF-IDF
→ 稀疏 × 稠密矩阵
→ ReLU
→ L2 Normalize
→ 若干稠密全连接层
→ 最终投影
→ L2 Normalize
→ dense embedding
1. 输入检查
assert bow_csr_matrix.shape[1] == self.layers[0].shape[1]
保证 TF-IDF 的词表维度和第一层输入维度一致。
例如：
bow_csr_matrix: (100, 2_000_000)
layer[0]: (92, 2_000_000)
表示 100 篇文档，每篇文档从约 200 万维投影到 92 维。
2. 分 batch 计算
for start in range(0, num_rows, batch_size):
避免一次处理所有文档造成内存过大。
bow_batch = bow_csr_matrix[start:end].astype(np.float32)
bow_batch 仍然是 CSR 稀疏矩阵，不会变成两百万维的稠密矩阵。
3. 第一层：稀疏转稠密
csr_matvecs_tiled_unrolled_8(
bow_batch.indptr,
bow_batch.indices,
bow_batch.data,
self.layers[0].T,
chunk_activations[0][:batch_size],
)
数学上等价于：
h1 = bow_batch @ self.layers[0].T
但它使用专门优化的 CSR kernel，只遍历 TF-IDF 中的非零特征。
对于一篇文档，某个输出维度为：
h1[j] = Σ TFIDF[k] × W0[j, k]
其中只计算文档实际包含的 ngram：
k ∈ 当前文档的非零特征
例如输入只有三个非零特征：
feature 10 = 0.2
feature 100 = 0.7
feature 500 = 0.4
第 j 个输出就是：
h1[j] =
0.2 × W[j, 10]
+ 0.7 × W[j, 100]
+ 0.4 × W[j, 500]
因为每个输出维度通常都会得到一个非零结果，所以这里完成了 sparse → dense 转换。
4. ReLU 和归一化
relu_inplace(h1)
normalize_inplace(h1)
ReLU：
ReLU(x) = max(0, x)
将负值变成零。
归一化通常是逐行 L2 Normalize：
h1 = h1 / ||h1||₂
让每篇文档的中间向量长度为 1。
5. 隐藏层
np.matmul(
chunk_activations[i - 1][:batch_size],
self.layers[i].T,
out=chunk_activations[i][:batch_size],
)
此时输入已经是 dense，因此使用普通矩阵乘法：
hi = h(i-1) @ Wi.T
hi = ReLU(hi)
hi = L2Normalize(hi)
每一层都执行：
线性投影 → ReLU → L2 归一化
这里没有看到 bias，因此只有权重矩阵。
6. 最后一层
np.matmul(
chunk_activations[-1][:batch_size],
self.layers[-1].T,
out=out[start:end],
)
normalize_inplace(out[start:end])
最后一层没有 ReLU：
embedding = h_last @ W_final.T
embedding = L2Normalize(embedding)
假设最后一层形状是：
(192, 3072)
最终输出：
(num_documents, 192)
每篇文档得到一个 192 维 dense embedding。
权重从哪里来
推理时不会调用 create() 随机初始化。权重来自 checkpoint：
state_dict[f"embedder.nn_layer_{i}"]
然后构造：
SparseToDenseEmbedder(layers=layers)
create() 的 Kaiming uniform 初始化主要用于新建或训练模型。
最终可以概括为：
两百万维稀疏 TF-IDF
→ 预训练稀疏投影层
→ 多层小型神经网络
→ 192 维归一化 dense embedding第一层 W 很大，因为它要把约 200 万维 ngram 词表映射到较小的隐藏维度。
如果第一层形状为：
W0.shape = (92, 2_000_000)
参数数量：
92 × 2,000,000 = 184,000,000
使用 float32 时：
184,000,000 × 4 bytes ≈ 736 MB
所以主要内存来自第一层大矩阵；CSR 优化节省的是输入和计算量，并不消除第一层权重本身的存储成本。
这个函数就是第一层的核心计算：
Y += A @ B
其中：
A = TF-IDF CSR 稀疏矩阵，shape=(batch_rows, vocab_size)
B = 第一层权重的转置，shape=(vocab_size, hidden_dim)
Y = 输出 dense 矩阵，shape=(batch_rows, hidden_dim)
调用处传入：
B_dense = self.layers[0].T
如果原始权重：
self.layers[0].shape = (92, 2_000_000)
那么：
B_dense.shape = (2_000_000, 92)
Y.shape = (batch_size, 92)
CSR 行遍历
对于第 i 篇文档：
for jj in range(A_indptr[i], A_indptr[i + 1]):
只遍历该文档非零的 TF-IDF 特征。
j = A_indices[jj]
a = A_data[jj]
含义：
j = ngram 的词表下标
a = 该 ngram 的 TF-IDF 权重
例如：
文档 i:
feature 2 → 0.2
feature 5 → 0.8
对应输出：
Y[i, :] += 0.2 × B[2, :]
Y[i, :] += 0.8 × B[5, :]
即：
Y[i, k] = Σ A[i, j] × B[j, k]
求和只遍历非零 j。
为什么每次处理 8 个输出维度
BLOCK_SIZE = 8
隐藏维度被拆成每组 8 个：
0～7
8～15
16～23
...
完整的 8 维块直接展开：
Y[i, k + 0] += ...
Y[i, k + 1] += ...
...
Y[i, k + 7] += ...
这样做可以：
减少 Python/Numba 循环控制开销
帮助编译器进行 SIMD 向量化
改善权重访问的缓存局部性
一次读取 a 后用于 8 个输出
如果输出维度是 92：
11 个完整块：11 × 8 = 88
最后一个剩余块：4
剩余的 4 维使用通用循环：
for k in range(88, 92):
Y[i, k] += a * B[j, k]
多线程
for i in numba.prange(n_row):
如果外层有：
@numba.njit(parallel=True, nogil=True)
不同文档行会由多个 CPU 线程并行计算。
这是线程安全的，因为不同线程分别写：
线程 1 → Y[0, :]
线程 2 → Y[1, :]
线程 3 → Y[2, :]
不会同时修改同一行。转成Dense特征的核心就是多个ngram特征, 根据W进行加权求和, 得到dense特征
多个 ngram 的 TF-IDF 特征
→ 分别取出 W 中对应的权重向量
→ 按 TF-IDF 权重加权求和
→ 得到第一层 dense 向量完整代码:
python
from __future__ import annotations
import logging
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Sequence, TypeVar, cast
import numba
import numpy as np
import pyarrow as pa
from numba.typed import Dict as NumbaTypedDict
from numpy.typing import NDArray
from scipy.sparse import csr_matrix
from tqdm.auto import tqdm
from luxical.ngrams import (
SpaceSavingNgramSummary,
bow_matrix_ngrams,
build_ngram_hash_to_idx_map,
)
from luxical.sparse_to_dense_neural_nets import SparseToDenseEmbedder
from luxical.tokenization import ArrowTokenizer
logger = logging.getLogger(__name__)
# Embedder serialization format version
EMBEDDER_FORMAT_VERSION = 1
@dataclass(frozen=True)
class Embedder:
tokenizer: ArrowTokenizer
recognized_ngrams: NDArray[np.int64]
ngram_hash_to_ngram_idx: NumbaTypedDict[np.int64, np.uint32]
idf_values: NDArray[np.float32]
bow_to_dense_embedder: SparseToDenseEmbedder
@property
def max_ngram_length(self) -> int:
return self.recognized_ngrams.shape[1]
@property
def embedding_dim(self) -> int:
return self.bow_to_dense_embedder.output_dim
def save(self, path: str | Path) -> None:
path = Path(path).resolve()
# Convert the ngram hash to ngram index map to a pair of numpy arrays.
ngram_idx_keys, ngram_idx_values = _unpack_int_dict(
self.ngram_hash_to_ngram_idx
)
# Name the individual neural network layers.
neural_layers = self.bow_to_dense_embedder.layers
layers_dict = {f"nn_layer_{i}": layer for i, layer in enumerate(neural_layers)}
# Store the tokenizer state JSON string as a numpy array of bytes.
tokenizer_bytes_arr = np.frombuffer(
self.tokenizer.to_str().encode("utf-8"), dtype=np.uint8
)
# Serialize the embedder state to a numpy file.
np.savez(
path,
tokenizer=tokenizer_bytes_arr,
version=np.int64(EMBEDDER_FORMAT_VERSION),
recognized_ngrams=self.recognized_ngrams,
ngram_hash_to_ngram_idx_keys=ngram_idx_keys,
ngram_hash_to_ngram_idx_values=ngram_idx_values,
idf_values=self.idf_values,
num_nn_layers=np.array([len(neural_layers)]),
**layers_dict,
allow_pickle=False,
)
@classmethod
def load(cls, path: str | Path) -> Embedder:
path = Path(path).resolve()
with np.load(path, allow_pickle=False) as npzfile:
# Ensure supported version.
version = npzfile["version"].item()
if version != EMBEDDER_FORMAT_VERSION:
raise NotImplementedError(
f"Unsupported embedder format version {version}. "
f"This code supports version {EMBEDDER_FORMAT_VERSION}."
)
# Deserialize the tokenizer.
tokenizer_bytes_arr = npzfile["tokenizer"]
tokenizer_str = tokenizer_bytes_arr.tobytes().decode("utf-8")
tokenizer = ArrowTokenizer(tokenizer_str)
# Deserialize the recognized ngrams.
recognized_ngrams = npzfile["recognized_ngrams"]
# Deserialize the ngram hash to ngram index map.
ngram_idx_keys = npzfile["ngram_hash_to_ngram_idx_keys"]
ngram_idx_values = npzfile["ngram_hash_to_ngram_idx_values"]
ngram_hash_to_ngram_idx = _pack_int_dict(ngram_idx_keys, ngram_idx_values)
# Deserialize the IDF values.
idf_values = npzfile["idf_values"]
# Deserialize the neural network layers into a `SparseToDenseEmbedder`.
num_nn_layers = npzfile["num_nn_layers"][0]
nn_layers = [npzfile[f"nn_layer_{i}"] for i in range(num_nn_layers)]
bow_to_dense_embedder = SparseToDenseEmbedder(layers=nn_layers)
return cls(
tokenizer=tokenizer,
recognized_ngrams=recognized_ngrams,
ngram_hash_to_ngram_idx=ngram_hash_to_ngram_idx,
bow_to_dense_embedder=bow_to_dense_embedder,
idf_values=idf_values,
)
def tokenize(
self, texts: pa.StringArray | Sequence[str]
) -> Sequence[NDArray[np.uint32]]:
token_pa = self.tokenizer.tokenize(pa.array(texts), add_special_tokens=False)
return token_pa.to_numpy(zero_copy_only=False)
def bow_from_tokens(
self, tokenized_docs: Iterable[NDArray[np.uint32]], progress_bar: bool = False
) -> csr_matrix:
return bow_matrix_ngrams(
tokenized_docs=tokenized_docs,
max_ngram_length=self.max_ngram_length,
ngram_hash_to_idx=self.ngram_hash_to_ngram_idx,
progress_bar=progress_bar,
)
def bow_from_texts(self, texts: Sequence[str]) -> csr_matrix:
tokenized_docs = self.tokenize(texts)
return self.bow_from_tokens(tokenized_docs, progress_bar=False)
def tfidf_from_bow(self, bow: csr_matrix) -> csr_matrix:
tfidf_data = np.empty_like(bow.data, dtype=np.float32)
_fast_tfidf_from_bow(
tfidf_data=tfidf_data,
bow_data=cast(NDArray[np.integer[Any]], bow.data),
bow_indices=bow.indices,
bow_indptr=bow.indptr,
idf_values=self.idf_values,
)
return csr_matrix((tfidf_data, bow.indices, bow.indptr), shape=bow.shape)
def replace_sparse_to_dense_embedder(
self, new_embedder: SparseToDenseEmbedder
) -> Embedder:
"""Replaces the neural network parameters as a copy operation."""
return Embedder(
tokenizer=self.tokenizer,
recognized_ngrams=self.recognized_ngrams,
ngram_hash_to_ngram_idx=self.ngram_hash_to_ngram_idx,
idf_values=self.idf_values,
bow_to_dense_embedder=new_embedder,
)
def __call__(
self,
texts: pa.StringArray | Sequence[str],
batch_size: int = 4096,
progress_bars: bool = False,
out: NDArray[np.float32] | None = None,
) -> NDArray[np.float32]:
"""Embed a sequence of texts into a dense vector space.
Internally batches data to avoid memory overhead. Optionally provides a
`tqdm` progress bar.
"""
n = len(texts)
out = np.empty((n, self.embedding_dim), dtype=np.float32)
with (
tqdm(
total=n,
desc="Embedding",
unit="text",
unit_scale=True,
disable=not progress_bars,
) as pbar,
tqdm(
unit="token", unit_scale=True, position=1, disable=not progress_bars
) as pbar_tokens,
):
for start in range(0, n, batch_size):
# Take the next batch of input texts.
end = min(start + batch_size, n)
batch_size = end - start
batch_text = texts[start:end]
# Tokenize the batch of texts.
tokenized_docs = self.tokenize(batch_text)
batch_size_in_tokens = np.vectorize(len)(tokenized_docs).sum()
# Convert the tokenized batch to a BoW matrix.
batch_bow = self.bow_from_tokens(tokenized_docs, progress_bar=False)
# TF-IDF.
batch_tfidf = self.tfidf_from_bow(batch_bow)
# Project the BoW matrix to a dense vector space.
self.bow_to_dense_embedder(batch_tfidf, out=out[start:end])
# Update the progress bars.
pbar.update(batch_size)
pbar_tokens.update(batch_size_in_tokens)
return out
def initialize_embedder_from_ngram_summary(
ngram_summary: SpaceSavingNgramSummary[np.int64],
tokenizer: ArrowTokenizer,
sparse_to_dense_embedder_dims: Sequence[int],
min_ngram_count_multiple: float = 10.0, # Minimum multiple of the lowest approximate count to keep an ngram.
max_vocabulary_size: int = 2_000_000, # Maximum number of ngrams to keep as features.
dtype: type[np.floating] = np.float32,
embedder_init_seed: int = 0,
) -> Embedder:
"""Construct an `Embedder` from the ngram statistics of a `SpaceSavingNgramSummary`.
Since the counts are approximate and the approximation is worse for the rarer ngrams,
let's throw out all ngrams with a count that is less than a certain multiple of the
lowest approximate count. We'll also limit the vocabulary size to a certain upper
bound.
"""
logger.info("Initializing embedder from features based on ngram summary")
min_approx_count = ngram_summary.approximate_counts.min()
keep_threshold = min_ngram_count_multiple * min_approx_count
is_above_threshold = ngram_summary.approximate_counts >= keep_threshold
fraction_above_threshold = is_above_threshold.mean()
num_above_threshold = int(is_above_threshold.sum())
vocabulary_size = min(max_vocabulary_size, num_above_threshold)
kept_counts = ngram_summary.approximate_counts[-vocabulary_size:]
kept_ngrams = ngram_summary.ngrams[-vocabulary_size:]
logger.info(
f"There are {num_above_threshold:,d} ngrams ({fraction_above_threshold:.2%}) "
f"with approximate count greater than {keep_threshold:,}, while the max "
f"vocabulary size is {max_vocabulary_size:,d} ngrams. This gives us a "
f"vocabulary of {vocabulary_size:,d} ngrams to track in our embedder."
)
# Pseudo-IDF (inverse log-scaled term frequency, ignoring the concept of "documents").
idf_values = (np.log(kept_counts.sum()) - np.log(kept_counts)).astype(dtype)
logger.info(
"Determined pseudo-IDF vector for the given vocabulary size of "
f"{vocabulary_size:,} ngrams with min value {idf_values.min():.2f} and "
f"max value {idf_values.max():.2f}."
)
# Also pull out the ngrams themselves and construct the ngram hash to index map for
# saving the full embedder state.
kept_ngrams = ngram_summary.ngrams[-vocabulary_size:]
ngram_hash_to_ngram_idx = NumbaTypedDict.empty(
numba.types.int64, numba.types.uint32
)
build_ngram_hash_to_idx_map(kept_ngrams, ngram_hash_to_ngram_idx)
# Randomly initialize the sparse to dense embedder.
sparse_to_dense_embedder = SparseToDenseEmbedder.create(
dims=[vocabulary_size] + list(sparse_to_dense_embedder_dims),
dtype=dtype,
seed=embedder_init_seed,
)
return Embedder(
tokenizer=tokenizer,
recognized_ngrams=kept_ngrams,
ngram_hash_to_ngram_idx=ngram_hash_to_ngram_idx,
idf_values=idf_values,
bow_to_dense_embedder=sparse_to_dense_embedder,
)
@numba.njit(error_model="numpy", parallel=True, nogil=True)
def _fast_tfidf_from_bow(
bow_data: NDArray[np.integer[Any]],
bow_indices: NDArray[np.integer[Any]],
bow_indptr: NDArray[np.integer[Any]],
idf_values: NDArray[np.floating[Any]],
tfidf_data: NDArray[np.floating[Any]],
l2_normalize: bool = True,
) -> None:
"""Fast and memory-efficient TF-IDF transformation of a Bag-of-Words CSR matrix."""
num_rows = bow_indptr.shape[0] - 1
for i in numba.prange(num_rows):
# Get the slice of data and indices for the current row.
start, end = bow_indptr[i], bow_indptr[i + 1]
row_tfidf = tfidf_data[start:end]
row_bow = bow_data[start:end]
row_indices = bow_indices[start:end]
# In the first pass over the row, apply IDF weighting and calculate the
# squared L2 norm of the resulting TF-IDF vector.
l2_norm_sq = 0.0
for j in range(row_bow.shape[0]):
weighted_val = row_bow[j] * idf_values[row_indices[j]]
row_tfidf[j] = weighted_val
l2_norm_sq += weighted_val * weighted_val
# In the second pass, normalize the row by the L2 norm.
if l2_normalize and l2_norm_sq > 0.0:
l2_norm = np.sqrt(l2_norm_sq)
row_tfidf /= l2_norm
# Serialization/deserialization helpers for numba typed dicts.
T_IntDictKey = TypeVar("T_IntDictKey", np.int32, np.int64)
T_IntDictValue = TypeVar("T_IntDictValue", np.int32, np.int64)
def _unpack_int_dict(
d: NumbaTypedDict[T_IntDictKey, T_IntDictValue],
) -> tuple[NDArray[T_IntDictKey], NDArray[T_IntDictValue]]:
key_dtype = numba.np.numpy_support.as_dtype(d._numba_type_.key_type) # type: ignore[attr-defined]
value_dtype = numba.np.numpy_support.as_dtype(d._numba_type_.value_type) # type: ignore[attr-defined]
keys = np.empty(len(d), dtype=key_dtype)
values = np.empty(len(d), dtype=value_dtype)
_unpack_int_dict_numba(d, keys, values) # type: ignore[arg-type]
return keys, values # type: ignore[return-value]
@numba.njit
def _unpack_int_dict_numba(
d: NumbaTypedDict[T_IntDictKey, T_IntDictValue],
keys: NDArray[T_IntDictKey],
values: NDArray[T_IntDictValue],
) -> None:
for i, (k, v) in enumerate(d.items()):
keys[i] = k
values[i] = v
def _pack_int_dict(
keys: NDArray[T_IntDictKey], values: NDArray[T_IntDictValue]
) -> NumbaTypedDict[T_IntDictKey, T_IntDictValue]:
assert len(keys) == len(values)
n = len(keys)
d = NumbaTypedDict.empty(
key_type=numba.from_dtype(keys.dtype),
value_type=numba.from_dtype(values.dtype),
n_keys=n,
)
_pack_int_dict_numba(d, keys, values)
return d
@numba.njit
def _pack_int_dict_numba(
d: NumbaTypedDict[T_IntDictKey, T_IntDictValue],
keys: NDArray[T_IntDictKey],
values: NDArray[T_IntDictValue],
) -> None:
for k, v in zip(keys, values):
d[k] = v