from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer from transformers.trainer_callback import EarlyStoppingCallback, ProgressCallback import torch from torch.utils.data import Dataset import numpy as np from sklearn.model_selection import train_test_split from seqeval.metrics import f1_score, precision_score, recall_score import random import os from ner_config import RepaymentNERConfig # 设置随机种子 def set_seed(seed): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) set_seed(RepaymentNERConfig.SEED) class RepaymentDataset(Dataset): def __init__(self, texts, labels, tokenizer, label_list): self.texts = texts self.labels = labels self.tokenizer = tokenizer self.label2id = {label: i for i, label in enumerate(label_list)} self.id2label = {i: label for i, label in enumerate(label_list)} self.encodings = self.tokenize_and_align_labels() def tokenize_and_align_labels(self): """分词并对齐标签""" tokenized_inputs = self.tokenizer( self.texts, # 直接传入文本列表 is_split_into_words=True, # 指示输入已经分词 truncation=True, padding=True, max_length=RepaymentNERConfig.MAX_LENGTH, return_offsets_mapping=True, return_tensors=None ) labels = [] for i, label_seq in enumerate(self.labels): word_ids = tokenized_inputs.word_ids(i) previous_word_idx = None label_ids = [] for word_idx in word_ids: if word_idx is None: # 特殊token，如[CLS], [SEP], [PAD] label_ids.append(-100) elif word_idx != previous_word_idx: # 新词的第一个token try: label_ids.append(self.label2id[label_seq[word_idx]]) except IndexError: print(f"错误：样本 {i} 的标签序列长度与文本不匹配") print(f"文本长度: {len(self.texts[i])}") print(f"标签长度: {len(label_seq)}") print(f"word_idx: {word_idx}") raise else: # 同一个词的后续token # 如果前一个token是实体的一部分，则使用相同的标签 if label_seq[word_idx-1].startswith("B-"): current_type = label_seq[word_idx-1][2:] label_ids.append(self.label2id[f"I-{current_type}"]) elif label_seq[word_idx-1].startswith("I-"): label_ids.append(self.label2id[label_seq[word_idx-1]]) else: label_ids.append(self.label2id["O"]) previous_word_idx = word_idx labels.append(label_ids) tokenized_inputs["labels"] = labels return tokenized_inputs def __getitem__(self, idx): return {key: torch.tensor(val[idx]) for key, val in self.encodings.items()} def __len__(self): return len(self.texts) def compute_metrics(p): predictions, labels = p predictions = np.argmax(predictions, axis=2) true_predictions = [ [RepaymentNERConfig.LABELS[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels) ] true_labels = [ [RepaymentNERConfig.LABELS[l] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels) ] results = { "overall_f1": f1_score(true_labels, true_predictions), "overall_precision": precision_score(true_labels, true_predictions), "overall_recall": recall_score(true_labels, true_predictions) } return results def train_repayment_model(texts, labels): # 加载预训练模型和分词器 tokenizer = AutoTokenizer.from_pretrained(RepaymentNERConfig.MODEL_NAME) model = AutoModelForTokenClassification.from_pretrained( RepaymentNERConfig.MODEL_NAME, num_labels=len(RepaymentNERConfig.LABELS), id2label={i: label for i, label in enumerate(RepaymentNERConfig.LABELS)}, label2id={label: i for i, label in enumerate(RepaymentNERConfig.LABELS)} ) # 划分训练集和验证集 train_texts, val_texts, train_labels, val_labels = train_test_split( texts, labels, test_size=RepaymentNERConfig.TEST_SIZE, random_state=RepaymentNERConfig.SEED ) # 创建数据集 train_dataset = RepaymentDataset(train_texts, train_labels, tokenizer, RepaymentNERConfig.LABELS) val_dataset = RepaymentDataset(val_texts, val_labels, tokenizer, RepaymentNERConfig.LABELS) # 训练参数 training_args = TrainingArguments( output_dir=RepaymentNERConfig.MODEL_PATH, num_train_epochs=RepaymentNERConfig.EPOCHS, per_device_train_batch_size=RepaymentNERConfig.BATCH_SIZE, per_device_eval_batch_size=RepaymentNERConfig.BATCH_SIZE * 2, # 评估时可以用更大的批次 warmup_ratio=RepaymentNERConfig.WARMUP_RATIO, weight_decay=RepaymentNERConfig.WEIGHT_DECAY, logging_dir=RepaymentNERConfig.LOG_PATH, logging_steps=RepaymentNERConfig.LOGGING_STEPS, evaluation_strategy="steps", eval_steps=RepaymentNERConfig.EVAL_STEPS, save_strategy="steps", save_steps=RepaymentNERConfig.EVAL_STEPS, save_total_limit=RepaymentNERConfig.SAVE_TOTAL_LIMIT, load_best_model_at_end=True, metric_for_best_model="overall_f1", greater_is_better=True, max_grad_norm=RepaymentNERConfig.MAX_GRAD_NORM, gradient_accumulation_steps=RepaymentNERConfig.GRADIENT_ACCUMULATION_STEPS, fp16=RepaymentNERConfig.FP16, dataloader_num_workers=RepaymentNERConfig.DATALOADER_NUM_WORKERS, dataloader_pin_memory=RepaymentNERConfig.DATALOADER_PIN_MEMORY, save_safetensors=True, optim="adamw_torch", disable_tqdm=False, report_to=["tensorboard"], group_by_length=True, # 相似长度的样本放在一起，减少padding length_column_name="length", remove_unused_columns=True, label_smoothing_factor=0.1, # 添加标签平滑 ) # 创建训练器 trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, callbacks=[ EarlyStoppingCallback( early_stopping_patience=RepaymentNERConfig.EARLY_STOPPING_PATIENCE, early_stopping_threshold=0.001 ), # 添加进度条回调 ProgressCallback() ] ) try: # 训练模型 print("\n开始训练模型...") train_result = trainer.train() # 打印训练结果 print("\n训练完成！") print(f"训练时长: {train_result.metrics['train_runtime']:.2f}秒") # 安全地获取和打印指标 metrics = train_result.metrics print("\n训练指标:") for key, value in metrics.items(): if isinstance(value, (int, float)): print(f"- {key}: {value:.4f}") # 最终评估 final_eval = trainer.evaluate() print("\n最终评估结果:") print(f"F1分数: {final_eval['eval_overall_f1']:.4f}") print(f"准确率: {final_eval['eval_overall_precision']:.4f}") print(f"召回率: {final_eval['eval_overall_recall']:.4f}") # 保存最佳模型 print("\n保存模型...") save_path = f"{RepaymentNERConfig.MODEL_PATH}/best_model" trainer.save_model(save_path) tokenizer.save_pretrained(save_path) print(f"模型已保存到: {save_path}") return model, tokenizer except Exception as e: print(f"\n训练过程中断: {str(e)}") # 尝试保存当前模型 try: save_path = f"{RepaymentNERConfig.MODEL_PATH}/interrupted_model" trainer.save_model(save_path) tokenizer.save_pretrained(save_path) print(f"已保存中断时的模型到: {save_path}") except Exception as save_error: print(f"保存中断模型失败: {str(save_error)}") raise def validate_labels(labels, valid_labels): """验证标签是否合法""" label_set = set() for seq in labels: label_set.update(seq) invalid_labels = label_set - set(valid_labels) if invalid_labels: raise ValueError(f"发现非法标签: {invalid_labels}") def clean_text(text: str) -> str: """清理文本中的特殊字符""" # 替换全角字符为半角 text = text.replace('￥', '¥') text = text.replace('，', ',') text = text.replace('。', '.') text = text.replace('：', ':') text = text.replace('（', '(') text = text.replace('）', ')') return text def preprocess_data(texts, labels): """预处理数据""" processed_texts = [] processed_labels = [] for i, (text, label_seq) in enumerate(zip(texts, labels)): if len(text) != len(label_seq): print(f"警告：样本 {i} 的文本和标签长度不匹配，已跳过") continue # 清理文本 cleaned_text = [clean_text(word) for word in text] # 处理金额标注 is_min_amount = False new_labels = [] for j, (word, label) in enumerate(zip(cleaned_text, label_seq)): if label.startswith("B-PICKUP_CODE"): # 检查是否是最低还款金额 context = ''.join(cleaned_text[max(0, j-5):j]) if any(kw in context for kw in RepaymentNERConfig.AMOUNT_CONFIG['min_amount_keywords']): is_min_amount = True new_labels.append("B-MIN_CODE") else: is_min_amount = False new_labels.append(label) elif label.startswith("I-PICKUP_CODE"): if is_min_amount: new_labels.append("I-MIN_CODE") else: new_labels.append(label) else: new_labels.append(label) processed_texts.append(cleaned_text) processed_labels.append(new_labels) return processed_texts, processed_labels def load_data(file_path): """加载并预处理数据""" texts = [] labels = [] current_words = [] current_labels = [] skip_url = False url_indicators = {'u', 'ur', 'url', 'http', 'https', 'www', 'com', 'cn'} def is_url_part(word): return (word.lower() in url_indicators or '.' in word or '/' in word or word.startswith('?')) with open(file_path, 'r', encoding='utf-8') as f: for line_num, line in enumerate(f, 1): line = line.strip() if not line: # 样本分隔符 if current_words: texts.append(current_words) labels.append(current_labels) current_words = [] current_labels = [] skip_url = False continue try: word, label = line.split(maxsplit=1) # URL处理逻辑 if is_url_part(word): skip_url = True elif word in ['】', '，', '。', '：']: skip_url = False # 标签验证 if not skip_url: if label not in RepaymentNERConfig.LABELS: print(f"警告：第{line_num}行发现非法标签 '{label}'，已跳过") continue current_words.append(word) current_labels.append(label) except Exception as e: print(f"错误：第{line_num}行处理失败 '{line}': {str(e)}") continue # 处理最后一个样本 if current_words: texts.append(current_words) labels.append(current_labels) return texts, labels def validate_dataset(texts, labels): """验证数据集的完整性和正确性""" stats = { "total_samples": len(texts), "total_tokens": sum(len(t) for t in texts), "entity_counts": {}, "avg_length": 0, "errors": [] } for i, (text, label_seq) in enumerate(zip(texts, labels)): # 长度检查 if len(text) != len(label_seq): stats["errors"].append(f"样本 {i}: 文本和标签长度不匹配") continue # 统计实体 current_entity = None for j, (word, label) in enumerate(zip(text, label_seq)): if label.startswith("B-"): entity_type = label[2:] stats["entity_counts"][entity_type] = stats["entity_counts"].get(entity_type, 0) + 1 current_entity = entity_type elif label.startswith("I-"): if not current_entity: stats["errors"].append(f"样本 {i}: 位置 {j} 的I-标签前没有对应的B-标签") elif label[2:] != current_entity: stats["errors"].append(f"样本 {i}: 位置 {j} 的I-标签类型与B-标签不匹配") else: current_entity = None stats["avg_length"] = stats["total_tokens"] / stats["total_samples"] if stats["total_samples"] > 0 else 0 return stats def resume_training(checkpoint_path): """从检查点恢复训练""" print(f"从检查点恢复训练: {checkpoint_path}") # 加载模型和分词器 tokenizer = AutoTokenizer.from_pretrained(checkpoint_path) model = AutoModelForTokenClassification.from_pretrained(checkpoint_path) # 重新加载数据 texts, labels = load_data(RepaymentNERConfig.DATA_PATH) texts, labels = preprocess_data(texts, labels) # 重新创建数据集 train_texts, val_texts, train_labels, val_labels = train_test_split( texts, labels, test_size=RepaymentNERConfig.TEST_SIZE, random_state=RepaymentNERConfig.SEED ) train_dataset = RepaymentDataset(train_texts, train_labels, tokenizer, RepaymentNERConfig.LABELS) val_dataset = RepaymentDataset(val_texts, val_labels, tokenizer, RepaymentNERConfig.LABELS) # 创建训练器并继续训练 training_args = TrainingArguments( output_dir=RepaymentNERConfig.MODEL_PATH, num_train_epochs=RepaymentNERConfig.EPOCHS, # ... 其他参数与train_repayment_model中相同 ... ) trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, callbacks=[EarlyStoppingCallback(early_stopping_patience=RepaymentNERConfig.EARLY_STOPPING_PATIENCE)] ) # 继续训练 trainer.train(resume_from_checkpoint=checkpoint_path) return model, tokenizer def main(): # 加载数据 print("正在加载数据...") texts, labels = load_data(RepaymentNERConfig.DATA_PATH) # 数据预处理 print("正在预处理数据...") texts, labels = preprocess_data(texts, labels) # 验证数据 print("验证数据集...") for i, (text, label_seq) in enumerate(zip(texts, labels)): if len(text) != len(label_seq): print(f"错误：样本 {i} 的文本和标签长度不匹配") print(f"文本({len(text)}): {text}") print(f"标签({len(label_seq)}): {label_seq}") return print(f"数据验证通过，共 {len(texts)} 个有效样本") # 验证数据集 print("正在验证数据集...") stats = validate_dataset(texts, labels) print("\n=== 数据集统计 ===") print(f"总样本数: {stats['total_samples']}") print(f"平均长度: {stats['avg_length']:.2f}") print("\n实体统计:") for entity_type, count in stats['entity_counts'].items(): print(f"- {entity_type}: {count}") if stats['errors']: print("\n发现以下问题:") for error in stats['errors']: print(f"- {error}") if input("是否继续训练? (y/n) ").lower() != 'y': return # 检查是否存在中断的模型 interrupted_model_path = f"{RepaymentNERConfig.MODEL_PATH}/interrupted_model" if os.path.exists(interrupted_model_path): print("\n发现中断的训练模型") if input("是否从中断处继续训练? (y/n) ").lower() == 'y': model, tokenizer = resume_training(interrupted_model_path) return # 正常训练流程 print("\n开始新的训练...") model, tokenizer = train_repayment_model(texts, labels) if __name__ == "__main__": main()