from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer from transformers.trainer_callback import EarlyStoppingCallback import torch from torch.utils.data import Dataset import numpy as np from sklearn.model_selection import train_test_split, KFold from seqeval.metrics import classification_report, f1_score, precision_score, recall_score import random import os from ner_config import NERConfig # 设置随机种子 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(NERConfig.SEED) # 更新标签列表，确保包含所有可能的标签 label_list = [ "O", "B-POST", "I-POST", "B-COMPANY", "I-COMPANY", "B-ADDRESS", "I-ADDRESS", "B-PICKUP_CODE", "I-PICKUP_CODE", # 添加取件码标签 "B-STATION", "I-STATION" # 添加站点标签 ] class NERDataset(Dataset): def __init__(self, texts, labels, tokenizer, label_list): self.texts = texts self.labels = labels self.tokenizer = tokenizer # 创建标签到ID的映射 self.label2id = {label: i for i, label in enumerate(label_list)} self.id2label = {i: label for i, label in enumerate(label_list)} # 打印标签映射信息，用于调试 print("标签映射:") for label, idx in self.label2id.items(): print(f"{label}: {idx}") # 对文本进行编码 self.encodings = self.tokenize_and_align_labels() def tokenize_and_align_labels(self): tokenized_inputs = self.tokenizer( [''.join(text) for text in self.texts], truncation=True, padding=True, max_length=NERConfig.MAX_LENGTH, return_offsets_mapping=True, return_tensors=None ) labels = [] for i, label in enumerate(self.labels): word_ids = tokenized_inputs.word_ids(i) previous_word_idx = None label_ids = [] current_entity = None for word_idx in word_ids: if word_idx is None: label_ids.append(-100) elif word_idx != previous_word_idx: # 新词开始 label_ids.append(self.label2id[label[word_idx]]) if label[word_idx].startswith("B-"): current_entity = label[word_idx][2:] elif label[word_idx] == "O": current_entity = None else: # 同一个词的后续token if current_entity: label_ids.append(self.label2id[f"I-{current_entity}"]) 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 load_data(file_path): texts, labels = [], [] current_words, current_labels = [], [] def clean_pickup_code_labels(words, labels): """清理取件码标注，确保不包含额外文字""" i = 0 while i < len(words): if labels[i].startswith("B-PICKUP_CODE"): # 找到取件码的结束位置 j = i + 1 while j < len(words) and ( labels[j].startswith("I-PICKUP_CODE") and (words[j].isalnum() or words[j] == "-") ): j += 1 # 检查并修正取件码序列 code_words = words[i:j] code_str = ''.join(code_words) # 如果发现无效字符或结束词，提前结束取件码 for k, word in enumerate(code_words): if not (word.isalnum() or word == "-") or word in ["的", "来", "到", "取", "件", "码"]: # 将后续标签改为O for m in range(i + k, j): labels[m] = "O" break i = j else: i += 1 return words, labels with open(file_path, 'r', encoding='utf-8') as f: for line in f: line = line.strip() if line: try: word, label = line.split(maxsplit=1) current_words.append(word) current_labels.append(label) except Exception as e: print(f"错误：处理行时出错: '{line}'") continue elif current_words: # 遇到空行且当前有数据 # 清理取件码标注 current_words, current_labels = clean_pickup_code_labels(current_words, current_labels) texts.append(current_words) labels.append(current_labels) current_words, current_labels = [], [] if current_words: # 处理最后一个样本 current_words, current_labels = clean_pickup_code_labels(current_words, current_labels) texts.append(current_words) labels.append(current_labels) return texts, labels def compute_metrics(p): """计算评估指标""" predictions, labels = p predictions = np.argmax(predictions, axis=2) # 移除特殊token的预测和标签 true_predictions = [ [NERConfig.LABELS[p] for (p, l) in zip(prediction, label) if l != -100] for prediction, label in zip(predictions, labels) ] true_labels = [ [NERConfig.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) } # 计算每个实体类型的指标 for entity_type in ["POST", "COMPANY", "ADDRESS", "PICKUP_CODE"]: # 将标签转换为二进制形式（是否属于当前实体类型） binary_preds = [] binary_labels = [] for pred_seq, label_seq in zip(true_predictions, true_labels): pred_binary = [] label_binary = [] for pred, label in zip(pred_seq, label_seq): # 检查标签是否属于当前实体类型 pred_is_entity = pred.endswith(entity_type) label_is_entity = label.endswith(entity_type) pred_binary.append(1 if pred_is_entity else 0) label_binary.append(1 if label_is_entity else 0) binary_preds.append(pred_binary) binary_labels.append(label_binary) # 计算当前实体类型的F1分数 try: entity_f1 = f1_score( sum(binary_labels, []), # 展平列表 sum(binary_preds, []), # 展平列表 average='binary' # 使用二进制评估 ) results[f"{entity_type}_f1"] = entity_f1 except Exception as e: print(f"计算{entity_type}的F1分数时出错: {str(e)}") results[f"{entity_type}_f1"] = 0.0 return results def augment_data(texts, labels): """数据增强""" augmented_texts = [] augmented_labels = [] for text, label in zip(texts, labels): # 原始数据 augmented_texts.append(text) augmented_labels.append(label) # 删除一些无关字符 if NERConfig.USE_DATA_AUGMENTATION: new_text = [] new_label = [] for t, l in zip(text, label): if l == "O" and random.random() < NERConfig.AUGMENTATION_RATIO: continue new_text.append(t) new_label.append(l) augmented_texts.append(new_text) augmented_labels.append(new_label) return augmented_texts, augmented_labels def train_with_different_seeds(texts, labels, n_seeds=2): """使用不同的随机种子进行多次训练""" best_model = None best_f1 = 0 results = [] # 创建输出目录 os.makedirs(NERConfig.MODEL_PATH, exist_ok=True) os.makedirs(NERConfig.LOG_PATH, exist_ok=True) for seed in range(n_seeds): print(f"\n=== 使用随机种子 {seed} 开始训练 ===") # 设置随机种子 set_seed(seed) # 加载分词器和模型 tokenizer = AutoTokenizer.from_pretrained(NERConfig.MODEL_NAME) model = AutoModelForTokenClassification.from_pretrained( NERConfig.MODEL_NAME, num_labels=len(NERConfig.LABELS), id2label={i: label for i, label in enumerate(NERConfig.LABELS)}, label2id={label: i for i, label in enumerate(NERConfig.LABELS)} ) # 创建数据集 train_texts, val_texts, train_labels, val_labels = train_test_split( texts, labels, test_size=NERConfig.TEST_SIZE, random_state=seed ) # 数据增强 if NERConfig.USE_DATA_AUGMENTATION: train_texts, train_labels = augment_data(train_texts, train_labels) train_dataset = NERDataset(train_texts, train_labels, tokenizer, NERConfig.LABELS) val_dataset = NERDataset(val_texts, val_labels, tokenizer, NERConfig.LABELS) # 训练参数 training_args = TrainingArguments( output_dir=f"{NERConfig.MODEL_PATH}/seed_{seed}", num_train_epochs=NERConfig.EPOCHS, per_device_train_batch_size=NERConfig.BATCH_SIZE, per_device_eval_batch_size=NERConfig.BATCH_SIZE, learning_rate=NERConfig.LEARNING_RATE, warmup_ratio=NERConfig.WARMUP_RATIO, weight_decay=NERConfig.WEIGHT_DECAY, gradient_accumulation_steps=NERConfig.GRADIENT_ACCUMULATION_STEPS, logging_steps=NERConfig.LOGGING_STEPS, save_total_limit=2, no_cuda=True, evaluation_strategy="steps", eval_steps=NERConfig.EVAL_STEPS, save_strategy="steps", save_steps=NERConfig.EVAL_STEPS, load_best_model_at_end=True, metric_for_best_model="overall_f1", greater_is_better=True, logging_dir=f"{NERConfig.LOG_PATH}/seed_{seed}", logging_first_step=True, report_to=["tensorboard"], ) # 训练器 trainer = Trainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=val_dataset, compute_metrics=compute_metrics, # 使用外部定义的compute_metrics callbacks=[EarlyStoppingCallback(early_stopping_patience=NERConfig.EARLY_STOPPING_PATIENCE)] ) # 训练模型 trainer.train() # 评估结果 eval_results = trainer.evaluate() results.append({ "seed": seed, "f1": eval_results["eval_overall_f1"], # 使用overall_f1而不是eval_f1 "metrics": eval_results }) # 如果是最好的模型，保存它 if eval_results["eval_overall_f1"] > best_f1: best_f1 = eval_results["eval_overall_f1"] best_model = (model, tokenizer) # 打印所有结果 print("\n=== 训练结果汇总 ===") for result in results: print(f"\n种子 {result['seed']} 的评估结果:") print(f"总体 F1: {result['f1']:.4f}") for key, value in result['metrics'].items(): if key.endswith('_f1'): print(f"{key}: {value:.4f}") # 保存最佳模型 if best_model: model, tokenizer = best_model save_path = f"{NERConfig.MODEL_PATH}/best_model" print(f"\n保存最佳模型到: {save_path}") model.save_pretrained(save_path) tokenizer.save_pretrained(save_path) return best_model, results def main(): # 加载数据 texts, labels = load_data(NERConfig.DATA_PATH) print(f"加载的数据集大小：{len(texts)}个样本") # 使用不同种子进行多次训练 best_model, results = train_with_different_seeds(texts, labels, n_seeds=NERConfig.N_SEEDS) if __name__ == "__main__": main()