import os
import time
import numpy as np
import pandas as pd
import pickle
import scipy.sparse
try:
_original_diags = scipy.sparse.diags
def _fixed_diags(diagonals, offsets=0, shape=None, format=None, dtype=None):
if isinstance(diagonals, np.ndarray) and diagonals.ndim == 1:
diagonals = [diagonals]
if np.isscalar(offsets):
offsets = [offsets]
return _original_diags(diagonals, offsets, shape, format, dtype)
scipy.sparse.diags = _fixed_diags
except Exception:
pass
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
from bertopic import BERTopic
from umap import UMAP
from hdbscan import HDBSCAN
from sklearn.feature_extraction.text import CountVectorizer, TfidfVectorizer
from sklearn.decomposition import NMF, LatentDirichletAllocation
from gensim import models, corpora
from sentence_transformers import SentenceTransformer
try:
from top2vec import Top2Vec
TOP2VEC_AVAILABLE = True
except ImportError:
TOP2VEC_AVAILABLE = False
def get_or_compute_embeddings(texts, save_path, model_name):
if os.path.exists(save_path):
print(f"[{time.strftime('%H:%M:%S')}] 加载向量文件...")
embeddings = np.load(save_path)
if len(embeddings) != len(texts):
print(f"数量不符，重新计算...")
else:
return embeddings
print(f"[{time.strftime('%H:%M:%S')}] 计算向量中...")
model = SentenceTransformer(model_name)
embeddings = model.encode(texts, show_progress_bar=True)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
np.save(save_path, embeddings)
print(f"向量已保存")
return embeddings
def save_aligned_corpus(texts, segmented_texts, save_dir):
with open(os.path.join(save_dir, 'aligned_full_texts.pkl'), 'wb') as f:
pickle.dump(texts, f)
with open(os.path.join(save_dir, 'aligned_segmented_texts.pkl'), 'wb') as f:
pickle.dump(segmented_texts, f)
def format_topic_words(words_iterable, n_top=10):
if words_iterable and isinstance(words_iterable[0], tuple):
return " ".join([w[0] for w in words_iterable][:n_top])
elif words_iterable:
return " ".join(words_iterable[:n_top])
return ""
def save_standard_results(model_name, topic_list, doc_topics, doc_probs, full_docs, save_dir):
pd.DataFrame(topic_list).to_csv(f"{save_dir}/{model_name}_topics.csv", index=False, encoding='utf-8-sig')
min_len = min(len(full_docs), len(doc_topics))
pd.DataFrame({
"Document": full_docs[:min_len],
"Topic": doc_topics[:min_len],
"Probability": doc_probs[:min_len]
}).to_csv(f"{save_dir}/{model_name}_docs.csv", index=False, encoding='utf-8-sig')
print(f"model_name} 结果已保存")
# ================= main =================
def run_pipeline():
RAW_TEXT_PATH = './data/文本.txt'
CUT_DATA_PATH = './data/切词.txt'
EMBEDDING_PATH = './data/embedding.npy'
RESULT_DIR = './result'
if not os.path.exists(RESULT_DIR): os.makedirs(RESULT_DIR)
print(f"[{time.strftime('%H:%M:%S')}] 1. 读取数据...")
try:
with open(RAW_TEXT_PATH, 'r', encoding='utf-8') as f:
full_texts = [line.strip() for line in f.readlines() if line.strip()]
with open(CUT_DATA_PATH, 'r', encoding='utf-8') as f:
segmented_docs = [line.strip() for line in f.readlines() if line.strip()]
except Exception as e:
print(f"❌ 错误：{e}")
return
min_len = min(len(full_texts), len(segmented_docs))
full_texts = full_texts[:min_len]
segmented_docs = segmented_docs[:min_len]
# 获取向量
embeddings = get_or_compute_embeddings(segmented_docs, EMBEDDING_PATH, "paraphrase-multilingual-MiniLM-L12-v2")
# 二次对齐
min_len = min(len(segmented_docs), len(embeddings))
segmented_docs = segmented_docs[:min_len]
full_texts = full_texts[:min_len]
embeddings = embeddings[:min_len]
save_aligned_corpus(full_texts, segmented_docs, RESULT_DIR)
stop_w = ['洛阳', '旅游', '文化', '我们', '一个']
# --- 1. BERTopic ---
print(f"\n[{time.strftime('%H:%M:%S')}] 正在训练 BERTopic...")
topic_model_bert = BERTopic(
umap_model=UMAP(n_neighbors=15, n_components=5, min_dist=0.0, metric='cosine', random_state=30),
hdbscan_model=HDBSCAN(min_cluster_size=10, min_samples=5, metric='euclidean', prediction_data=True),
vectorizer_model=CountVectorizer(token_pattern=r"(?u)[^ ]+", stop_words=stop_w),
verbose=False
)
topics_bert, probs_bert = topic_model_bert.fit_transform(segmented_docs, embeddings=embeddings)
bert_info = topic_model_bert.get_topic_info()
bert_std_topics = []
for _, row in bert_info.iterrows():
if row['Topic'] == -1: continue
bert_std_topics.append({
"Topic": row['Topic'], "Count": row['Count'], "Words": format_topic_words(row['Representation'])
})
if isinstance(probs_bert, np.ndarray) and probs_bert.ndim > 1:
final_probs = probs_bert.max(axis=1)
else:
final_probs = probs_bert if hasattr(probs_bert, '__len__') else [1.0]*len(topics_bert)
save_standard_results("BERTopic", bert_std_topics, topics_bert, final_probs, full_texts, RESULT_DIR)
# 确定主题数
num_topics_fixed = len(bert_std_topics) if len(bert_std_topics) >= 2 else 5
print(f"💡 后续模型统一主题数: {num_topics_fixed}")
# --- 2. LDA (Gensim) ---
print(f"[{time.strftime('%H:%M:%S')}] 正在训练 LDA...")
texts_split = [doc.split() for doc in segmented_docs]
dictionary = corpora.Dictionary(texts_split)
corpus = [dictionary.doc2bow(text) for text in texts_split]
lda = models.LdaModel(corpus=corpus, id2word=dictionary, num_topics=num_topics_fixed, random_state=42)
lda_res = [(sorted(lda[doc], key=lambda x: -x[1])[0]) if doc else (-1, 0) for doc in corpus]
lda_topics, lda_probs = zip(*lda_res)
lda_std = [{"Topic": i, "Words": format_topic_words(lda.show_topic(i, topn=15))} for i in range(num_topics_fixed)]
save_standard_results("LDA", lda_std, lda_topics, lda_probs, full_texts, RESULT_DIR)
# 准备 TF-IDF 矩阵 (给 pLSA 和 NMF 用)
tfidf = TfidfVectorizer(token_pattern=r"(?u)[^ ]+", stop_words=stop_w, max_features=2000)
X_tfidf = tfidf.fit_transform(segmented_docs)
feats = tfidf.get_feature_names_out()
# --- 3. pLSA (使用 Sklearn LDA 近似) ---
print(f"[{time.strftime('%H:%M:%S')}] 正在训练 pLSA...")
# 使用 batch 学习方法的 LDA 在 Sklearn 中常被用作 pLSA 的替代
plsa = LatentDirichletAllocation(n_components=num_topics_fixed, learning_method='batch', random_state=42)
W_plsa = plsa.fit_transform(X_tfidf)
plsa_std = [{"Topic": i, "Words": " ".join([feats[j] for j in comp.argsort()[:-11:-1]])} for i, comp in enumerate(plsa.components_)]
save_standard_results("pLSA", plsa_std, W_plsa.argmax(axis=1), W_plsa.max(axis=1), full_texts, RESULT_DIR)
# --- 4. NMF ---
print(f"[{time.strftime('%H:%M:%S')}] 正在训练 NMF...")
nmf = NMF(n_components=num_topics_fixed, random_state=42, init='nndsvda')
W_nmf = nmf.fit_transform(X_tfidf)
nmf_std = [{"Topic": i, "Words": " ".join([feats[j] for j in comp.argsort()[:-11:-1]])} for i, comp in enumerate(nmf.components_)]
save_standard_results("NMF", nmf_std, W_nmf.argmax(axis=1), W_nmf.max(axis=1), full_texts, RESULT_DIR)
# --- 5. Top2Vec ---
if TOP2VEC_AVAILABLE:
print(f"[{time.strftime('%H:%M:%S')}] 正在训练 Top2Vec...")
try:
t2v = Top2Vec(documents=segmented_docs, speed="deep-attention", workers=4, embedding_model='paraphrase-multilingual-MiniLM-L12-v2')
num_t = t2v.get_num_topics()
words, scores, nums = t2v.get_topics(num_t)
t2v_std = [{"Topic": n, "Words": format_topic_words(list(w))} for n, w in zip(nums, words)]
save_standard_results("Top2Vec", t2v_std, t2v.doc_top, t2v.doc_dist, full_texts, RESULT_DIR)
except Exception as e:
print(f"❌ Top2Vec 失败: {e}")
print(f"\n🎉 全部完成！")
if __name__ == "__main__":
run_pipeline()