准备工作
pip install gensim jieba markdown scikit-learn
gensim: Doc2Vec 的核心库。jieba: 中文分词库，处理中文文档必不可少。markdown: 用于将Markdown文档转换为纯文本，以便Doc2Vec处理。scikit-learn: 用于计算文档相似度。
步骤
- 数据预处理: 加载Markdown文档并进行预处理，包括分词、去除停用词等。
- 训练Doc2Vec模型: 使用预处理后的文档训练Doc2Vec模型，生成文档向量。
- 相似度计算: 使用训练好的模型计算文档之间的相似度。
- 结果展示: 展示相似度最高的文档及其相似度分数。
引入依赖
import os
import jieba
import re
from gensim.models.doc2vec import Doc2Vec, TaggedDocument
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report
import markdown
根据输入的Markdown文档目录加载和预处理文档
目录如图所示
# --- 1. 数据收集与预处理 ---
def load_and_preprocess_markdown_documents(data_dir):
"""
加载指定目录下所有Markdown文档，进行预处理（去除Markdown语法、分词）。
假设每个子文件夹代表一个类别。
"""
documents = []
labels = []
doc_id_counter = 0
for category_name in os.listdir(data_dir):
category_path = os.path.join(data_dir, category_name)
if os.path.isdir(category_path):
print(f"Processing category: {category_name}")
for filename in os.listdir(category_path):
if filename.endswith(".md"):
filepath = os.path.join(category_path, filename)
with open(filepath, 'r', encoding='utf-8') as f:
md_content = f.read()
# 将Markdown转换为纯文本
html = markdown.markdown(md_content)
plain_text = re.sub('<[^<]+?>', '', html).strip() # 去除HTML标签
# 简单的文本清理，去除特殊字符和数字，并转换为小写
plain_text = re.sub(r'[^\w\s]', '', plain_text).lower()
plain_text = re.sub(r'\d+', '', plain_text)
# 中英文分词
words = []
if re.search(r'[\u4e00-\u9fa5]', plain_text): # 包含中文字符
words = list(jieba.cut(plain_text))
else: # 纯英文
words = plain_text.split()
# 过滤空字符串
words = [word for word in words if word.strip()]
if words:
documents.append(TaggedDocument(words=words, tags=[f'DOC_{doc_id_counter}']))
labels.append(category_name)
doc_id_counter += 1
return documents, labels
模型训练
# --- 2. Doc2Vec模型训练 ---
def train_doc2vec_model(tagged_documents, vector_size=100, window=5, min_count=1, epochs=20):
"""
训练Doc2Vec模型。
"""
print("Training Doc2Vec model...")
model = Doc2Vec(vector_size=vector_size, window=window, min_count=min_count, workers=4, epochs=epochs)
model.build_vocab(tagged_documents)
model.train(tagged_documents, total_examples=model.corpus_count, epochs=model.epochs)
print("Doc2Vec model training complete.")
return model
- vector_size: 文档向量的维度。维度越高，捕捉语义信息的能力越强，存储的信息越多。
- window: 上下文窗口大小。窗口越大，考虑的上下文越广泛。窗口越小，更关注局部词语关系。对于短文档，使用小窗口（2-5），长文档使用大窗口（5-15）。
- min_count: 忽略出现次数小于min_count的单词。
- epochs: 训练轮数。
生成文档向量并训练分类器
# --- 3. 文档向量生成 ---
def get_document_vectors(model, tagged_documents):
"""
从训练好的Doc2Vec模型中获取每个文档的向量。
"""
doc_vectors = []
for doc in tagged_documents:
doc_vectors.append(model.dv[doc.tags[0]]) # 获取通过tag训练的向量
return doc_vectors
# --- 4. 分类器训练 ---
def train_classifier(X_train, y_train):
"""
训练一个分类器。
"""
print("Training classifier...")
classifier = LogisticRegression(max_iter=1000) # 增加max_iter以避免收敛警告
classifier.fit(X_train, y_train)
print("Classifier training complete.")
return classifier
开始训练并评估
# --- 主训练流程 ---
if __name__ == "__main__":
training_data_dir = 'data/training_docs' # 请替换为您的实际数据目录
print("--- Step 1: Loading and Preprocessing Documents ---")
all_tagged_documents, all_labels = load_and_preprocess_markdown_documents(training_data_dir)
if not all_tagged_documents:
print("No documents found or processed. Please check your data directory.")
else:
# 划分训练集和测试集 (用于分类器，而非Doc2Vec本身)
# Doc2Vec训练需要所有文档，但分类器需要独立测试集来评估性能
X_train_docs, X_test_docs, y_train, y_test = train_test_split(
all_tagged_documents, all_labels, test_size=0.4, random_state=42, stratify=all_labels
)
print(f"Total documents: {len(all_tagged_documents)}")
print(f"Training documents for classifier: {len(X_train_docs)}")
print(f"Testing documents for classifier: {len(X_test_docs)}")
# Doc2Vec训练使用所有文档，因为它学习的是文档的表示
print("\n--- Step 2: Training Doc2Vec Model ---")
doc2vec_model = train_doc2vec_model(all_tagged_documents, vector_size=150, epochs=30) # 增加向量维度和训练轮数
# 保存Doc2Vec模型，以便后续加载和使用
doc2vec_model_path = 'doc2vec_programming_docs.model'
doc2vec_model.save(doc2vec_model_path)
print(f"Doc2Vec model saved to {doc2vec_model_path}")
print("\n--- Step 3: Generating Document Vectors for Classifier Training ---")
# 为分类器训练集生成向量
X_train_vectors = get_document_vectors(doc2vec_model, X_train_docs)
print("\n--- Step 4: Training Classifier ---")
classifier = train_classifier(X_train_vectors, y_train)
# 保存分类器模型
import joblib
classifier_model_path = 'document_classifier.joblib'
joblib.dump(classifier, classifier_model_path)
print(f"Classifier model saved to {classifier_model_path}")
print("\n--- Step 5: Evaluating Classifier on Test Set ---")
X_test_vectors = get_document_vectors(doc2vec_model, X_test_docs)
y_pred = classifier.predict(X_test_vectors)
print("\nClassification Report:")
print(classification_report(y_test, y_pred))
运行结果如下：
- 整体表现 ：
- 准确率(accuracy)为77%，说明模型在测试集上能正确分类约3/4的文档
- 宏平均(macro avg)和加权平均(weighted avg)的F1分数都在0.77左右，表明模型整体表现尚可但仍有提升空间
- 各类别表现 ：
- Python类别 表现最好：
- 精确率(precision)=1.00：所有被预测为Python的文档确实都是Python
- 召回率(recall)=0.80：能识别出80%的Python文档
- C#类别 表现中等：
- 精确率和召回率都是0.75
- Java类别 表现相对最弱：
- 精确率0.60：预测为Java的文档中有40%其实是其他类别
- 召回率0.75：能识别出75%的Java文档
- 改进建议 ：
- 增加Java类别的训练样本，特别是容易与其他类别混淆的文档
- 检查C#和Java文档的内容相似性，可能需要调整特征提取方式
- 可以尝试调整Doc2Vec的参数（如vector_size）或使用更复杂的分类器
- 当前限制 ：
- 测试集样本量仍较小（共13个样本）
- 某些类别的样本可能不够代表性
- 结果可能会有较大波动
根据AI的分析，可以看出我们仍需要更多的样本数据来改进这个简单的编程语言分类模型。
