| | |
|---|
| | | # train_train_ner.py |
|---|
| | | 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 |
|---|
| | | from seqeval.metrics import f1_score, precision_score, recall_score |
|---|
| | | import random |
|---|
| | | import os |
|---|
| | | import re |
|---|
| | | from ner_config import TrainNERConfig |
|---|
| | | |
|---|
| | | # 设置随机种子 |
|---|
| | | 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(TrainNERConfig.SEED) |
|---|
| | | |
|---|
| | | 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=TrainNERConfig.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_trips_labels(words, labels): |
|---|
| | | """清理车次标注，确保格式正确""" |
|---|
| | | i = 0 |
|---|
| | | while i < len(words): |
|---|
| | | if labels[i].startswith("B-TRIPS"): # 修改标签名 |
|---|
| | | # 找到车次的结束位置 |
|---|
| | | j = i + 1 |
|---|
| | | while j < len(words) and labels[j].startswith("I-TRIPS"): # 修改标签名 |
|---|
| | | j += 1 |
|---|
| | | |
|---|
| | | # 检查并修正车次序列 |
|---|
| | | train_words = words[i:j] |
|---|
| | | train_str = ''.join(train_words) |
|---|
| | | |
|---|
| | | # 检查格式是否符合车次规范 |
|---|
| | | valid_patterns = [ |
|---|
| | | re.compile(r'^[GDCZTKY]\d{1,2}$'), |
|---|
| | | re.compile(r'^[GDCZTKY]\d{1,2}/\d{1,2}$'), |
|---|
| | | re.compile(r'^[GDCZTKY]\d{1,2}-\d{1,2}$'), |
|---|
| | | re.compile(r'^\d{1,4}$'), |
|---|
| | | re.compile(r'^[A-Z]\d{1,4}$') |
|---|
| | | ] |
|---|
| | | |
|---|
| | | is_valid = any(pattern.match(train_str) for pattern in valid_patterns) |
|---|
| | | if not is_valid: |
|---|
| | | # 将格式不正确的标签改为O |
|---|
| | | for k in range(i, j): |
|---|
| | | labels[k] = "O" |
|---|
| | | |
|---|
| | | 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_trips_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_trips_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 = [ |
|---|
| | | [TrainNERConfig.LABELS[p] for (p, l) in zip(prediction, label) if l != -100] |
|---|
| | | for prediction, label in zip(predictions, labels) |
|---|
| | | ] |
|---|
| | | true_labels = [ |
|---|
| | | [TrainNERConfig.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 ["COMPANY","TRIPS", "START", "END", "DATE", "TIME", "SEAT", "NAME"]: |
|---|
| | | # 将标签转换为二进制形式 |
|---|
| | | 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) |
|---|
| | | |
|---|
| | | # 删除一些无关字符 |
|---|
| | | new_text = [] |
|---|
| | | new_label = [] |
|---|
| | | for t, l in zip(text, label): |
|---|
| | | if l == "O" and random.random() < 0.3: |
|---|
| | | 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 main(): |
|---|
| | | # 加载数据 |
|---|
| | | texts, labels = load_data(TrainNERConfig.DATA_PATH) |
|---|
| | | print(f"加载的数据集大小：{len(texts)}个样本") |
|---|
| | | |
|---|
| | | # 划分数据集 |
|---|
| | | train_texts, val_texts, train_labels, val_labels = train_test_split( |
|---|
| | | texts, labels, test_size=TrainNERConfig.TEST_SIZE, random_state=TrainNERConfig.SEED |
|---|
| | | ) |
|---|
| | | |
|---|
| | | # 数据增强 |
|---|
| | | train_texts, train_labels = augment_data(train_texts, train_labels) |
|---|
| | | |
|---|
| | | # 加载分词器和模型 |
|---|
| | | tokenizer = AutoTokenizer.from_pretrained(TrainNERConfig.MODEL_NAME) |
|---|
| | | model = AutoModelForTokenClassification.from_pretrained( |
|---|
| | | TrainNERConfig.MODEL_NAME, |
|---|
| | | num_labels=len(TrainNERConfig.LABELS), |
|---|
| | | id2label={i: label for i, label in enumerate(TrainNERConfig.LABELS)}, |
|---|
| | | label2id={label: i for i, label in enumerate(TrainNERConfig.LABELS)} |
|---|
| | | ) |
|---|
| | | |
|---|
| | | # 创建数据集 |
|---|
| | | train_dataset = NERDataset(train_texts, train_labels, tokenizer, TrainNERConfig.LABELS) |
|---|
| | | val_dataset = NERDataset(val_texts, val_labels, tokenizer, TrainNERConfig.LABELS) |
|---|
| | | |
|---|
| | | # 训练参数 |
|---|
| | | training_args = TrainingArguments( |
|---|
| | | output_dir=TrainNERConfig.MODEL_PATH, |
|---|
| | | num_train_epochs=TrainNERConfig.EPOCHS, |
|---|
| | | per_device_train_batch_size=TrainNERConfig.BATCH_SIZE, |
|---|
| | | per_device_eval_batch_size=TrainNERConfig.BATCH_SIZE, |
|---|
| | | learning_rate=TrainNERConfig.LEARNING_RATE, |
|---|
| | | warmup_ratio=TrainNERConfig.WARMUP_RATIO, |
|---|
| | | weight_decay=TrainNERConfig.WEIGHT_DECAY, |
|---|
| | | gradient_accumulation_steps=TrainNERConfig.GRADIENT_ACCUMULATION_STEPS |
|---|
| | | ) |
|---|
| | | |
|---|
| | | trainer = Trainer( |
|---|
| | | model=model, |
|---|
| | | args=training_args, |
|---|
| | | train_dataset=train_dataset, |
|---|
| | | eval_dataset=val_dataset, |
|---|
| | | compute_metrics=compute_metrics |
|---|
| | | ) |
|---|
| | | |
|---|
| | | trainer.train() |
|---|
| | | # 评估结果 |
|---|
| | | eval_results = trainer.evaluate() |
|---|
| | | print("\n评估结果:") |
|---|
| | | for key, value in eval_results.items(): |
|---|
| | | print(f"{key}: {value:.4f}") |
|---|
| | | |
|---|
| | | # 保存最终模型 |
|---|
| | | model.save_pretrained(f"{TrainNERConfig.MODEL_PATH}/best_model") |
|---|
| | | tokenizer.save_pretrained(f"{TrainNERConfig.MODEL_PATH}/best_model") |
|---|
| | | |
|---|
| | | if __name__ == "__main__": |
|---|
| | | main() |
|---|
