from collections import Counter import numpy as np import pandas as pd import torch import matplotlib.pyplot as plt from sklearn.metrics import accuracy_score, classification_report from sklearn.model_selection import train_test_split from sklearn.preprocessing import StandardScaler from sklearn.svm import SVC from torch.utils.data import DataLoader, Dataset from tqdm import tqdm from transformers import AutoTokenizer, BertModel import joblib from sklearn.metrics import confusion_matrix import seaborn as sns # 1. ====================== 配置参数 ====================== MODEL_PATH = r'D:\pythonProject5\bert-base-chinese' BATCH_SIZE = 64 MAX_LENGTH = 128 SAVE_DIR = r'D:\pythonProject5\BSVMC_model' DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # 2. ====================== 数据加载与划分 ====================== def load_data(file_path): """加载并预处理数据""" df = pd.read_csv(file_path).dropna(subset=['text', 'label']) texts = df['text'].astype(str).str.strip().tolist() labels = df['label'].astype(int).tolist() return texts, labels # 加载原始数据 texts, labels = load_data("train3.csv") # 第一次拆分：分出测试集（20%） train_val_texts, test_texts, train_val_labels, test_labels = train_test_split( texts, labels, test_size=0.2, stratify=labels, random_state=42 ) # 第二次拆分：分出训练集（70%）和验证集（30% of 80% = 24%） train_texts, val_texts, train_labels, val_labels = train_test_split( train_val_texts, train_val_labels, test_size=0.3, # 0.3 * 0.8 = 24% of original stratify=train_val_labels, random_state=42 ) # 3. ====================== 文本编码 ====================== tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) def encode_texts(texts): return tokenizer( texts, truncation=True, padding="max_length", max_length=MAX_LENGTH, return_tensors="pt" ) # 编码所有数据集 train_encodings = encode_texts(train_texts) val_encodings = encode_texts(val_texts) test_encodings = encode_texts(test_texts) # 4. ====================== 数据集类 ====================== class TextDataset(Dataset): def __init__(self, encodings, labels): self.encodings = encodings self.labels = labels def __getitem__(self, idx): return { 'input_ids': self.encodings['input_ids'][idx], 'attention_mask': self.encodings['attention_mask'][idx], 'labels': torch.tensor(self.labels[idx]) } def __len__(self): return len(self.labels) # 创建所有数据集加载器 train_dataset = TextDataset(train_encodings, train_labels) val_dataset = TextDataset(val_encodings, val_labels) test_dataset = TextDataset(test_encodings, test_labels) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False) test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False) # 5. ====================== 特征提取 ====================== def extract_features(bert_model, dataloader): """使用BERT提取CLS特征""" bert_model.eval() all_features = [] all_labels = [] with torch.no_grad(): for batch in tqdm(dataloader, desc="提取特征"): inputs = {k: v.to(DEVICE) for k, v in batch.items() if k != 'labels'} outputs = bert_model(**inputs) features = outputs.last_hidden_state[:, 0, :].cpu().numpy() all_features.append(features) all_labels.append(batch['labels'].numpy()) return np.vstack(all_features), np.concatenate(all_labels) # 加载并冻结BERT模型 bert_model = BertModel.from_pretrained(MODEL_PATH).to(DEVICE) for param in bert_model.parameters(): param.requires_grad = False # 提取所有特征 print("\n" + "=" * 30 + " 特征提取阶段 " + "=" * 30) train_features, train_labels = extract_features(bert_model, train_loader) val_features, val_labels = extract_features(bert_model, val_loader) test_features, test_labels = extract_features(bert_model, test_loader) # 6. ====================== 特征预处理 ====================== scaler = StandardScaler() train_features = scaler.fit_transform(train_features) # 只在训练集上fit val_features = scaler.transform(val_features) test_features = scaler.transform(test_features) # 7. ====================== 训练SVM ====================== print("\n" + "=" * 30 + " 训练SVM模型 " + "=" * 30) svm_model = SVC( kernel='rbf', C=1.0, gamma='scale', probability=True, random_state=42 ) svm_model.fit(train_features, train_labels) # 8. ====================== 评估模型 ====================== def evaluate(features, labels, model, dataset_name): preds = model.predict(features) acc = accuracy_score(labels, preds) print(f"\n[{dataset_name}] 评估结果：") print(f"准确率：{acc:.4f}") print(classification_report(labels, preds, digits=4)) return preds print("\n训练集评估：") _ = evaluate(train_features, train_labels, svm_model, "训练集") print("\n验证集评估：") val_preds = evaluate(val_features, val_labels, svm_model, "验证集") print("\n测试集评估：") test_preds = evaluate(test_features, test_labels, svm_model, "测试集") # 9. ====================== 保存模型 ====================== def save_pipeline(): """保存完整模型管道""" # 创建保存目录 import os os.makedirs(SAVE_DIR, exist_ok=True) # 保存BERT相关 bert_model.save_pretrained(SAVE_DIR) tokenizer.save_pretrained(SAVE_DIR) # 保存SVM和预处理 joblib.dump(svm_model, f"{SAVE_DIR}/svm_model.pkl") joblib.dump(scaler, f"{SAVE_DIR}/scaler.pkl") # 保存标签映射（假设标签为0: "中性", 1: "正面", 2: "负面"） label_map = {0: "中性", 1: "正面", 2: "负面"} joblib.dump(label_map, f"{SAVE_DIR}/label_map.pkl") print(f"\n模型已保存至 {SAVE_DIR} 目录") save_pipeline() # 10. ===================== 可视化 ====================== plt.figure(figsize=(15, 5)) # 决策值分布 plt.subplot(1, 2, 1) plt.plot(svm_model.decision_function(train_features[:100]), 'o', alpha=0.5) plt.title("训练集前100样本决策值分布") plt.xlabel("样本索引") plt.ylabel("决策值") # 生成混淆矩阵 cm = confusion_matrix(y_true=test_labels, y_pred=test_preds) # 可视化 plt.figure(figsize=(10, 7)) sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=['0', '1', '2'], yticklabels=['0', '1', '2']) plt.xlabel('Predicted label') plt.ylabel('True label') plt.title('confusion matrix') plt.show() # 准确率对比 plt.subplot(1, 2, 2) accuracies = [ accuracy_score(train_labels, svm_model.predict(train_features)), accuracy_score(val_labels, val_preds), accuracy_score(test_labels, test_preds) ] labels = ['train', 'Validation', 'test'] plt.bar(labels, accuracies, color=['blue', 'orange', 'green']) plt.ylim(0, 1) plt.title("Comparison of accuracy rates for each dataset") plt.ylabel("Accuracy rate") plt.tight_layout() plt.show()画一下我的模型架构图