# 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()
|
