""" BERT合规性分类器训练脚本 使用Hugging Face Transformers训练序列分类模型 """ import os import json from pathlib import Path from typing import List, Dict import math import torch from torch.utils.data import Dataset from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, TrainingArguments, Trainer, DataCollatorWithPadding ) from sklearn.metrics import accuracy_score, precision_recall_fscore_support class ComplianceDataset(Dataset): """合规性分类数据集""" def __init__(self, data: List[Dict], tokenizer, label2id: Dict): self.data = data self.tokenizer = tokenizer self.label2id = label2id def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] encoding = self.tokenizer( item['text'], truncation=True, max_length=128, padding=False # DataCollator会处理padding ) encoding['label'] = self.label2id[item['label']] return encoding class BERTClassifierTrainer: """BERT分类器训练器""" def __init__( self, model_name: str, train_data_path: str, output_dir: str, val_data_path: str = None ): """ Args: model_name: 预训练模型名称 train_data_path: 训练数据路径 output_dir: 输出目录 val_data_path: 验证数据路径 """ self.model_name = model_name self.train_data_path = train_data_path self.val_data_path = val_data_path self.output_dir = Path(output_dir) self.output_dir.mkdir(parents=True, exist_ok=True) # 加载标签映射 label_mapping_path = Path(train_data_path).parent / 'label_mapping.json' if label_mapping_path.exists(): with open(label_mapping_path, 'r', encoding='utf-8') as f: label_mapping = json.load(f) self.label2id = label_mapping['label2id'] self.id2label = label_mapping['id2label'] self.num_labels = label_mapping['num_labels'] else: # 从数据中提取标签 self._extract_labels_from_data() print(f"标签数量: {self.num_labels}") print(f"标签映射: {self.label2id}") # 加载tokenizer print(f"\n加载tokenizer: {model_name}") self.tokenizer = AutoTokenizer.from_pretrained(model_name) # 检测GPU self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print(f"使用设备: {self.device}") if torch.cuda.is_available(): print(f"GPU: {torch.cuda.get_device_name(0)}") def _extract_labels_from_data(self): """从数据中提取标签""" print("\n从训练数据中提取标签...") # 加载训练数据 train_data = self._load_data(self.train_data_path) # 提取唯一标签 unique_labels = set(item['label'] for item in train_data) unique_labels = sorted(unique_labels) self.label2id = {label: idx for idx, label in enumerate(unique_labels)} self.id2label = {idx: label for label, idx in self.label2id.items()} self.num_labels = len(unique_labels) # 保存标签映射 label_mapping = { 'label2id': self.label2id, 'id2label': self.id2label, 'num_labels': self.num_labels } mapping_path = self.output_dir / 'label_mapping.json' with open(mapping_path, 'w', encoding='utf-8') as f: json.dump(label_mapping, f, ensure_ascii=False, indent=2) def _load_data(self, data_path: str) -> List[Dict]: """加载数据""" print(f"加载数据: {data_path}") data_path = Path(data_path) data = [] if data_path.suffix == '.jsonl': with open(data_path, 'r', encoding='utf-8') as f: for line in f: data.append(json.loads(line.strip())) else: with open(data_path, 'r', encoding='utf-8') as f: data = json.load(f) print(f"加载样本数: {len(data)}") return data def train( self, num_epochs: int = 3, batch_size: int = 16, learning_rate: float = 2e-5, warmup_steps: int = 100, weight_decay: float = 0.01 ): """训练模型""" print("\n" + "="*50) print("开始训练") print("="*50) # 加载训练数据 train_data = self._load_data(self.train_data_path) train_dataset = ComplianceDataset( train_data, self.tokenizer, self.label2id ) # 加载验证数据 eval_dataset = None if self.val_data_path and Path(self.val_data_path).exists(): val_data = self._load_data(self.val_data_path) eval_dataset = ComplianceDataset( val_data, self.tokenizer, self.label2id ) # 加载模型 print(f"\n加载模型: {self.model_name}") model = AutoModelForSequenceClassification.from_pretrained( self.model_name, num_labels=self.num_labels, id2label=self.id2label, label2id=self.label2id ) # 训练参数 total_steps = math.ceil(len(train_dataset) / batch_size) * num_epochs training_args = TrainingArguments( output_dir=str(self.output_dir), num_train_epochs=num_epochs, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, warmup_steps=warmup_steps, weight_decay=weight_decay, learning_rate=learning_rate, logging_dir=str(self.output_dir / 'logs'), logging_steps=10, eval_strategy="steps" if eval_dataset else "no", eval_steps=50 if eval_dataset else None, save_strategy="steps", save_steps=50, save_total_limit=3, load_best_model_at_end=True if eval_dataset else False, metric_for_best_model="f1" if eval_dataset else None, greater_is_better=True, report_to=None, # 不使用wandb/tensorboard ) print(f"\n训练配置:") print(f" 训练样本: {len(train_dataset)}") print(f" 批次大小: {batch_size}") print(f" 训练轮数: {num_epochs}") print(f" 总步数: {total_steps}") print(f" 学习率: {learning_rate}") # 创建Trainer trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, data_collator=DataCollatorWithPadding(self.tokenizer), compute_metrics=self._compute_metrics if eval_dataset else None, ) # 开始训练 print("\n开始训练...") trainer.train() # 保存最终模型 final_model_path = self.output_dir / 'final_model' trainer.save_model(str(final_model_path)) self.tokenizer.save_pretrained(str(final_model_path)) print(f"\n✓ 最终模型已保存: {final_model_path}") return trainer def _compute_metrics(self, eval_pred): """计算评估指标""" predictions, labels = eval_pred predictions = predictions.argmax(axis=-1) precision, recall, f1, _ = precision_recall_fscore_support( labels, predictions, average='weighted' ) accuracy = accuracy_score(labels, predictions) return { 'accuracy': accuracy, 'f1': f1, 'precision': precision, 'recall': recall } def evaluate(self, test_data_path: str = None): """评估模型""" if test_data_path is None: test_data_path = self.val_data_path if test_data_path is None or not Path(test_data_path).exists(): print("\n未提供测试数据，跳过评估") return print(f"\n评估模型: {test_data_path}") # 加载模型 model_path = self.output_dir / 'final_model' model = AutoModelForSequenceClassification.from_pretrained(model_path) # 加载数据 test_data = self._load_data(test_data_path) test_dataset = ComplianceDataset( test_data, self.tokenizer, self.label2id ) # 创建Trainer trainer = Trainer( model=model, data_collator=DataCollatorWithPadding(self.tokenizer) ) # 评估 results = trainer.evaluate(test_dataset) print(f"\n评估结果:") for key, value in results.items(): if 'eval_' in key: print(f" {key.replace('eval_', '')}: {value:.4f}") return results def main(): """主函数""" import argparse parser = argparse.ArgumentParser(description='训练BERT合规性分类器') parser.add_argument( '--model_name', type=str, default='hfl/chinese-bert-wwm-ext', help='预训练模型名称' ) parser.add_argument( '--train_data', type=str, required=True, help='训练数据路径' ) parser.add_argument( '--val_data', type=str, default=None, help='验证数据路径' ) parser.add_argument( '--output_dir', type=str, default='models/bert-compliance', help='输出目录' ) parser.add_argument( '--num_epochs', type=int, default=3, help='训练轮数' ) parser.add_argument( '--batch_size', type=int, default=16, help='批次大小' ) parser.add_argument( '--learning_rate', type=float, default=2e-5, help='学习率' ) args = parser.parse_args() # 设置随机种子 import random import numpy as np random.seed(42) np.random.seed(42) torch.manual_seed(42) # 创建训练器 trainer = BERTClassifierTrainer( model_name=args.model_name, train_data_path=args.train_data, output_dir=args.output_dir, val_data_path=args.val_data ) # 训练 trainer.train( num_epochs=args.num_epochs, batch_size=args.batch_size, learning_rate=args.learning_rate ) # 评估 if args.val_data: trainer.evaluate() print("\n✓ 训练完成!") print(f"\n模型已保存到: {args.output_dir}") print("\n使用方法:") print(f" from transformers import AutoModelForSequenceClassification") print(f" model = AutoModelForSequenceClassification.from_pretrained('{args.output_dir}/final_model')") if __name__ == '__main__': main()