Python案例如何用Doc2Vec文档向量

本文目录导读：

1. 安装必要的库
2. 基础Doc2Vec案例
3. 完整的相似度计算案例
4. 高级应用：文档聚类
5. 完整实用类封装
6. 模型保存和加载
7. 关键参数说明

我来详细介绍如何使用Doc2Vec来生成文档向量,包括完整的案例代码。

安装必要的库

pip install gensim nltk numpy

基础Doc2Vec案例

1 准备数据并训练模型

import gensim
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from nltk.tokenize import word_tokenize
import nltk
import numpy as np
# 下载必要的NLTK数据
nltk.download('punkt')
nltk.download('punkt_tab')
# 准备示例文档数据
documents = [
    "机器学习是人工智能的重要分支",
    "深度学习使用神经网络进行建模",
    "自然语言处理可以处理文本数据",
    "计算机视觉主要处理图像和视频",
    "强化学习通过奖励机制进行训练",
    "Python是数据科学中最流行的编程语言",
    "TensorFlow和PyTorch是常用的深度学习框架",
    "数据预处理是机器学习项目中的关键步骤"
]
# 创建TaggedDocument格式的数据
tagged_data = []
for i, doc in enumerate(documents):
    # 分词并创建TaggedDocument
    words = word_tokenize(doc)
    tagged_data.append(TaggedDocument(words=words, tags=[i]))
print(f"准备 {len(tagged_data)} 个带标签的文档")

2 训练Doc2Vec模型

# 创建和训练Doc2Vec模型
model = Doc2Vec(
    vector_size=100,      # 向量维度
    window=3,             # 上下文窗口大小
    min_count=1,          # 忽略词频低于此值的词
    dm=1,                 # 使用DM模型（1）或DBOW模型（0）
    workers=4,            # 并行训练线程数
    epochs=20,            # 训练迭代次数
    seed=42               # 随机种子
)
# 构建词汇表
model.build_vocab(tagged_data)
# 训练模型
model.train(
    tagged_data,
    total_examples=model.corpus_count,
    epochs=model.epochs
)
print("模型训练完成")
print(f"词汇表大小: {len(model.wv.key_to_index)}")

3 获取文档向量

# 获取训练好的文档向量
print("=== 获取文档向量 ===")
for i, doc in enumerate(documents):
    doc_vector = model.dv[i]  # 或 model.docvecs[i]
    print(f"文档 {i}: {doc[:20]}...")
    print(f"  向量维度: {doc_vector.shape}")
    print(f"  前5个值: {doc_vector[:5]}")
    print()
# 获取新文档的向量（推断）
new_doc = "深度学习在自然语言处理中应用广泛"
new_doc_words = word_tokenize(new_doc)
inferred_vector = model.infer_vector(new_doc_words)
print("=== 新文档向量 ===")
print(f"新文档: {new_doc}")
print(f"推断向量维度: {inferred_vector.shape}")
print(f"前5个值: {inferred_vector[:5]}")

完整的相似度计算案例

# 继续使用上面的模型
def find_similar_documents(model, query, top_n=3):
    """找到与查询最相似的文档"""
    # 对查询进行分词
    query_words = word_tokenize(query)
    # 推断查询向量
    query_vector = model.infer_vector(query_words)
    # 计算与所有文档的相似度
    similarities = model.dv.most_similar([query_vector], topn=top_n)
    return similarities
# 示例查询
queries = [
    "神经网络如何处理数据",
    "编程语言在AI中的应用",
    "图像识别的技术"
]
print("=== 文档相似度搜索 ===")
for query in queries:
    print(f"\n查询: {query}")
    results = find_similar_documents(model, query, top_n=3)
    for doc_idx, similarity in results:
        original_doc = documents[doc_idx]
        print(f"  文档 {doc_idx}: {original_doc}")
        print(f"  相似度: {similarity:.4f}")
# 计算两个文档之间的相似度
doc1_idx = 0
doc2_idx = 1
similarity = np.dot(model.dv[doc1_idx], model.dv[doc2_idx])
print(f"\n文档0和文档1的余弦相似度: {similarity:.4f}")

高级应用：文档聚类

from sklearn.cluster import KMeans
from sklearn.manifold import TSNE
import matplotlib.pyplot as plt
# 获取所有文档向量
doc_vectors = [model.dv[i] for i in range(len(documents))]
doc_vectors = np.array(doc_vectors)
# 执行K-means聚类
n_clusters = 2
kmeans = KMeans(n_clusters=n_clusters, random_state=42)
cluster_labels = kmeans.fit_predict(doc_vectors)
# 显示聚类结果
print("=== 文档聚类结果 ===")
for i, (doc, label) in enumerate(zip(documents, cluster_labels)):
    print(f"文档 {i}: {doc[:30]}... -> 聚类 {label}")
# 使用t-SNE可视化
tsne = TSNE(n_components=2, random_state=42)
doc_vectors_2d = tsne.fit_transform(doc_vectors)
# 绘制可视化
plt.figure(figsize=(10, 8))
scatter = plt.scatter(
    doc_vectors_2d[:, 0], 
    doc_vectors_2d[:, 1], 
    c=cluster_labels, 
    cmap='viridis', 
    s=100
)
# 添加文档标签
for i, doc in enumerate(documents):
    plt.annotate(
        f"Doc {i}", 
        (doc_vectors_2d[i, 0], doc_vectors_2d[i, 1]),
        xytext=(5, 5), 
        textcoords='offset points',
        fontsize=8
    )
'文档向量t-SNE可视化')
plt.colorbar(scatter, label='Cluster')
plt.tight_layout()
plt.show()

完整实用类封装

class Doc2VecProcessor:
    """Doc2Vec处理器类"""
    def __init__(self, vector_size=100, window=3, dm=1, epochs=20):
        self.vector_size = vector_size
        self.window = window
        self.dm = dm
        self.epochs = epochs
        self.model = None
        self.documents = None
    def fit(self, documents, progress_per=100):
        """训练Doc2Vec模型"""
        self.documents = documents
        # 准备数据
        tagged_data = []
        for i, doc in enumerate(documents):
            words = word_tokenize(doc)
            tagged_data.append(TaggedDocument(words=words, tags=[i]))
        # 创建模型
        self.model = Doc2Vec(
            vector_size=self.vector_size,
            window=self.window,
            min_count=1,
            dm=self.dm,
            workers=4,
            epochs=self.epochs
        )
        # 训练
        self.model.build_vocab(tagged_data)
        self.model.train(
            tagged_data,
            total_examples=self.model.corpus_count,
            epochs=self.model.epochs
        )
        return self
    def get_document_vector(self, doc_id):
        """获取指定文档的向量"""
        if self.model is None:
            raise ValueError("模型未训练")
        return self.model.dv[doc_id]
    def infer_vector(self, text):
        """推断新文本的向量"""
        if self.model is None:
            raise ValueError("模型未训练")
        words = word_tokenize(text)
        return self.model.infer_vector(words)
    def find_similar(self, query, top_n=5):
        """查找相似文档"""
        if self.model is None or self.documents is None:
            raise ValueError("模型未训练")
        query_vector = self.infer_vector(query)
        similarities = self.model.dv.most_similar([query_vector], topn=top_n)
        results = []
        for doc_idx, similarity in similarities:
            results.append({
                'document': self.documents[doc_idx],
                'index': doc_idx,
                'similarity': similarity
            })
        return results
    def save(self, filepath):
        """保存模型"""
        if self.model:
            self.model.save(filepath)
    def load(self, filepath):
        """加载模型"""
        self.model = Doc2Vec.load(filepath)
# 使用示例
processor = Doc2VecProcessor()
processor.fit(documents)
# 测试相似度搜索
query_text = "神经网络和深度学习"
results = processor.find_similar(query_text, top_n=3)
print(f"\n查询: {query_text}")
print("搜索结果:")
for result in results:
    print(f"  文档: {result['document']}")
    print(f"  相似度: {result['similarity']:.4f}")
    print()

模型保存和加载

# 保存模型
model.save("doc2vec_model.bin")
print("模型已保存")
# 加载模型
loaded_model = Doc2Vec.load("doc2vec_model.bin")
print("模型已加载")
# 使用加载的模型进行推断
new_doc = "人工智能技术发展迅速"
new_vector = loaded_model.infer_vector(word_tokenize(new_doc))
print(f"加载模型的推断向量维度: {new_vector.shape}")

关键参数说明

这个完整的Doc2Vec案例涵盖了从数据准备、模型训练到应用的所有步骤，可以直接用于你的项目中。