如何高效将Sklearn TF-IDF稀疏词嵌入矩阵转为Gensim KeyedVectors

嘿，我刚好处理过类似的场景，针对你用Sklearn TF-IDF得到的大型稀疏词嵌入矩阵，要转成Gensim KeyedVectors的需求，给你几个实用的方案，兼顾高效存储和易用性：

方案1：直接从稀疏矩阵创建KeyedVectors实例（无需提前保存） #

如果你的内存足够容纳转成稠密后的向量，这个方法最直接，一步到位：

from gensim.models import KeyedVectors
import scipy.sparse

# 假设你已经准备好以下两个变量：
# vocab_list: 词汇列表，顺序和稀疏矩阵的行严格对应
# sparse_embeddings: scipy.sparse.csr_matrix，形状为 (len(vocab_list), embedding_dim)

# 初始化KeyedVectors对象，指定向量维度
kv = KeyedVectors(vector_size=sparse_embeddings.shape[1])

# 将稀疏矩阵转为稠密数组（注意：超大型矩阵可能会占用大量内存，谨慎使用）
dense_embeddings = sparse_embeddings.toarray()

# 把词汇和对应的向量绑定到KeyedVectors中
kv.add_vectors(vocab_list, dense_embeddings)

# 现在就可以像使用普通KeyedVectors一样操作了
print(kv["example_word"])

方案2：高效存储稀疏矩阵，加载后构建KeyedVectors #

对于超大型稀疏矩阵，直接转稠密会撑爆内存，这时候先把矩阵和词汇表分开高效存储，之后再分批加载构建KeyedVectors是更稳妥的选择：

第一步：保存稀疏矩阵和词汇表 #

import scipy.sparse
import pickle

# 用scipy的npz格式保存稀疏矩阵，这是专门为稀疏矩阵优化的存储格式，磁盘占用极小
scipy.sparse.save_npz("tfidf_embeddings.npz", sparse_embeddings)

# 用pickle保存词汇表（也可以存成txt，pickle更方便直接加载为列表）
with open("vocab_list.pkl", "wb") as f:
    pickle.dump(vocab_list, f)

第二步：加载并分批构建KeyedVectors #

from gensim.models import KeyedVectors
import scipy.sparse
import pickle

# 加载保存的稀疏矩阵和词汇表
sparse_embeddings = scipy.sparse.load_npz("tfidf_embeddings.npz")
with open("vocab_list.pkl", "rb") as f:
    vocab_list = pickle.load(f)

# 初始化KeyedVectors
kv = KeyedVectors(vector_size=sparse_embeddings.shape[1])

# 分批次转稠密并添加向量，避免一次性占用过多内存
batch_size = 1000  # 可以根据你的内存情况调整批次大小
for i in range(0, len(vocab_list), batch_size):
    batch_vocab = vocab_list[i:i+batch_size]
    batch_dense = sparse_embeddings[i:i+batch_size].toarray()
    kv.add_vectors(batch_vocab, batch_dense)

# 验证一下功能是否正常
print(kv.most_similar("target_word"))

额外技巧：自定义稀疏版KeyedVectors（可选） #

如果你的场景大部分时候只需要查询单个词的向量，不想把整个矩阵转成稠密，可以自定义一个包装类，内部保留稀疏矩阵和词汇表，只在查询时转成稠密向量，最大化节省内存：

from sklearn.metrics.pairwise import cosine_similarity

class SparseKeyedVectors:
    def __init__(self, vocab_list, sparse_embeddings):
        self.vocab_map = {word: idx for idx, word in enumerate(vocab_list)}
        self.sparse_matrix = sparse_embeddings
        self.vector_size = sparse_embeddings.shape[1]
    
    def __getitem__(self, word):
        idx = self.vocab_map.get(word)
        if idx is None:
            raise KeyError(f"Word '{word}' not in vocabulary")
        # 只把当前查询的词向量转成稠密
        return self.sparse_matrix[idx].toarray().flatten()
    
    def most_similar(self, positive=[], negative=[], topn=10):
        # 计算查询向量，再用稀疏矩阵的余弦相似度找出最相似的词
        query_vec = sum(self[word] for word in positive) - sum(self[word] for word in negative)
        sim_scores = cosine_similarity(query_vec.reshape(1, -1), self.sparse_matrix)[0]
        # 排序取topN
        top_indices = sim_scores.argsort()[::-1][:topn]
        return [(vocab_list[idx], sim_scores[idx]) for idx in top_indices]

# 使用示例
sparse_kv = SparseKeyedVectors(vocab_list, sparse_embeddings)
print(sparse_kv["example_word"])
print(sparse_kv.most_similar(positive=["example_word"]))

这个自定义类完美模拟了KeyedVectors的核心接口，同时保留了稀疏矩阵的存储优势，适合内存紧张的场景。

内容的提问来源于stack exchange，提问作者Chi Tang