Gensim中KeyedVector替代Doc2Vec.infer_vector()获取文档向量的方法

嘿，我来帮你理清这个问题！首先得明确一个关键点：KeyedVectors本身没有直接对应Doc2Vec中infer_vector的方法——因为Doc2Vec的文档向量推断依赖于训练时的完整模型架构（比如PV-DM/PV-DBOW的隐层计算逻辑），而KeyedVectors只是一个轻量的向量存储工具，只保存了预训练好的词/实体向量，丢掉了模型的训练推断逻辑。

不过根据你的场景（要生成文档向量写入Annoy），有两种靠谱的方案，看你是否还保留着原Doc2Vec模型：

方案1：保留原Doc2Vec模型（最贴合原效果，推荐） #

如果你的KeyedVectors是从原Doc2Vec模型导出的（比如通过doc2vec_model.wv），那最稳妥的方式是继续用原Doc2Vec模型的infer_vector生成文档向量，再存入Annoy。毕竟Doc2Vec的推断逻辑是专门为文档向量设计的，能考虑到文档的上下文序列信息，效果会比简单聚合词向量好很多。

示例代码：

from gensim.models.doc2vec import Doc2Vec
from annoy import AnnoyIndex

# 加载你的原Doc2Vec模型
doc2vec_model = Doc2Vec.load("your_trained_doc2vec.model")

# 封装推断文档向量的函数（参数可以和你之前的训练设置对齐）
def infer_doc_vector(doc_tokens):
    return doc2vec_model.infer_vector(
        doc_tokens, 
        alpha=0.025,  # 学习率，和训练时一致或微调
        steps=50      # 迭代步数，步数越多结果越稳定
    )

# 生成文档向量并构建Annoy索引
vector_dim = doc2vec_model.vector_size
annoy_index = AnnoyIndex(vector_dim, 'angular')  # 距离度量按需选择

# 假设你有分词后的文档列表docs_tokens
for doc_id, tokens in enumerate(docs_tokens):
    doc_vec = infer_doc_vector(tokens)
    annoy_index.add_item(doc_id, doc_vec)

# 构建索引，树的数量越多检索越准但占用资源越多
annoy_index.build(n_trees=10)
annoy_index.save("document_vectors.ann")

方案2：用KeyedVectors的词向量聚合生成文档向量 #

如果已经没有原Doc2Vec模型，只能基于KeyedVectors来生成文档向量，那可以通过聚合文档中所有词的向量来实现。常见的聚合方式有两种：

方式A：简单平均（最易用） #

直接对文档中所有存在于KeyedVectors的词向量取平均，实现简单且效果不差：

from gensim.models import KeyedVectors
from annoy import AnnoyIndex
import numpy as np

# 加载KeyedVectors
kv_model = KeyedVectors.load("your_keyedvectors.kv")
vector_dim = kv_model.vector_size

def get_avg_doc_vector(doc_tokens):
    # 过滤掉不在词向量模型中的词
    valid_vecs = [kv_model[token] for token in doc_tokens if token in kv_model]
    if not valid_vecs:
        # 无匹配词时返回全零向量，也可以根据需求调整默认值
        return np.zeros(vector_dim)
    return np.mean(valid_vecs, axis=0)

# 写入Annoy的逻辑和方案1一致
annoy_index = AnnoyIndex(vector_dim, 'angular')
for doc_id, tokens in enumerate(docs_tokens):
    doc_vec = get_avg_doc_vector(tokens)
    annoy_index.add_item(doc_id, doc_vec)

annoy_index.build(n_trees=10)
annoy_index.save("avg_doc_vectors.ann")

方式B：TF-IDF加权平均（更贴合词的重要性） #

如果想让文档中更重要的词（TF-IDF权重高的词）对向量影响更大，可以先训练TF-IDF模型，再做加权聚合：

from gensim.models import KeyedVectors
from gensim.corpora import Dictionary
from gensim.models import TfidfModel
from annoy import AnnoyIndex
import numpy as np

# 加载KeyedVectors
kv_model = KeyedVectors.load("your_keyedvectors.kv")
vector_dim = kv_model.vector_size

# 用你的文档语料构建字典和TF-IDF模型
doc_dict = Dictionary(docs_tokens)
corpus = [doc_dict.doc2bow(tokens) for tokens in docs_tokens]
tfidf_model = TfidfModel(corpus)

def get_tfidf_weighted_vector(doc_tokens):
    bow = doc_dict.doc2bow(doc_tokens)
    tfidf_weights = tfidf_model[bow]
    
    weighted_sum = np.zeros(vector_dim)
    total_weight = 0.0
    for token_id, weight in tfidf_weights:
        token = doc_dict[token_id]
        if token in kv_model:
            weighted_sum += kv_model[token] * weight
            total_weight += weight
    
    if total_weight == 0:
        return np.zeros(vector_dim)
    return weighted_sum / total_weight

# 构建Annoy索引
annoy_index = AnnoyIndex(vector_dim, 'angular')
for doc_id, tokens in enumerate(docs_tokens):
    doc_vec = get_tfidf_weighted_vector(tokens)
    annoy_index.add_item(doc_id, doc_vec)

annoy_index.build(n_trees=10)
annoy_index.save("tfidf_weighted_doc_vectors.ann")

最后提醒 #

• 如果你希望文档向量的效果和原Doc2Vec的infer_vector尽可能一致，务必保留原Doc2Vec模型——词向量聚合的方式无法复现Doc2Vec对文档序列信息的建模能力，效果会打折扣。
• 要是你只是想让词向量查询更轻量，完全可以用KeyedVectors处理词向量查询，文档向量还是用原Doc2Vec模型生成，这样兼顾轻量性和效果。

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