Scikit-learn 中 LLM 嵌入、TF-IDF 和词袋模型：哪个效果更好？

from sklearn.datasets import fetch_20newsgroups
from sklearn.model_selection import train_test_split

# 仅选择两个类别
CATEGORIES = ['rec.sport.baseball', 'sci.electronics']

# 加载数据
groups = fetch_20newsgroups(subset='all', categories=CATEGORIES, remove=('headers', 'footers', 'quotes'))

X_train, X_test, y_train, y_test = train_test_split(groups.data, groups.target, test_size=0.3, random_state=42)

print(f"训练样本数: {len(X_train)}")
print(f"测试样本数: {len(X_test)}")

from sklearn.feature_extraction.text import CountVectorizer

# 初始化并训练 CountVectorizer
vectorizer_bow = CountVectorizer()
# 注意：fit_transform 仅在训练集上进行
X_train_bow = vectorizer_bow.fit_transform(X_train)
X_test_bow = vectorizer_bow.transform(X_test)

print(f"BoW 特征形状: {X_train_bow.shape}")

from sklearn.feature_extraction.text import TfidfVectorizer

# 初始化并训练 TfidfVectorizer
vectorizer_tfidf = TfidfVectorizer()
X_train_tfidf = vectorizer_tfidf.fit_transform(X_train)
X_test_tfidf = vectorizer_tfidf.transform(X_test)

print(f"TF-IDF 特征形状: {X_train_tfidf.shape}")

# 假设已安装 sentence-transformers
from sentence_transformers import SentenceTransformer

# 选择一个轻量级且性能良好的模型
MODEL_NAME = 'all-MiniLM-L6-v2'

# 加载模型
embedding_model = SentenceTransformer(MODEL_NAME)

# 生成嵌入
X_train_emb = embedding_model.encode(X_train, convert_to_numpy=True)
X_test_emb = embedding_model.encode(X_test, convert_to_numpy=True)

print(f"LLM 嵌入特征形状: {X_train_emb.shape}")

from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score

# 用于存储结果
results = {}

# ----------------------
# 评估 BoW
# ----------------------
print("\n评估 BoW...")
classifier_bow = LogisticRegression(max_iter=1000)
classifier_bow.fit(X_train_bow, y_train)
y_pred_bow = classifier_bow.predict(X_test_bow)
accuracy_bow = accuracy_score(y_test, y_pred_bow)
results['BoW'] = accuracy_bow
print(f"BoW 准确率: {accuracy_bow:.4f}")

# ----------------------
# 评估 TF-IDF
# ----------------------
print("\n评估 TF-IDF...")
classifier_tfidf = LogisticRegression(max_iter=1000)
classifier_tfidf.fit(X_train_tfidf, y_train)
y_pred_tfidf = classifier_tfidf.predict(X_test_tfidf)
accuracy_tfidf = accuracy_score(y_test, y_pred_tfidf)
results['TF-IDF'] = accuracy_tfidf
print(f"TF-IDF 准确率: {accuracy_tfidf:.4f}")

# ----------------------
# 评估 LLM 嵌入
# ----------------------
print("\n评估 LLM 嵌入...")
# 嵌入特征通常需要C-SVM或线性模型，但我们使用相同的逻辑回归进行直接比较
# 对于嵌入，我们可能需要稍微调整正则化或增加迭代次数
classifier_emb = LogisticRegression(max_iter=2000)
classifier_emb.fit(X_train_emb, y_train)
y_pred_emb = classifier_emb.predict(X_test_emb)
accuracy_emb = accuracy_score(y_test, y_pred_emb)
results['LLM Embeddings'] = accuracy_emb
print(f"LLM 嵌入 准确率: {accuracy_emb:.4f}")

# ----------------------
# 最终对比
# ----------------------
print("\n--- 最终结果对比 ---")

import pandas as pd

results_df = pd.DataFrame(list(results.items()), columns=['Method', 'Accuracy'])
results_df = results_df.sort_values(by='Accuracy', ascending=False)

print(results_df.to_string(index=False))