You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

107 lines
2.9 KiB

import json
from datasets import Dataset
from transformers import (
AutoTokenizer,
AutoModelForTokenClassification,
TrainingArguments,
Trainer,
DataCollatorForTokenClassification
)
# 加载数据
def load_data(file_path):
with open(file_path, "r", encoding="utf-8") as f:
data = json.load(f)
return data
# 数据预处理
def tokenize_and_align_labels(examples, tokenizer):
tokenized_inputs = tokenizer(examples["text"], truncation=True, padding=True, is_split_into_words=True)
labels = []
for i, label in enumerate(examples["labels"]):
word_ids = tokenized_inputs.word_ids(batch_index=i)
label_ids = []
for word_idx in word_ids:
if word_idx is None:
label_ids.append(-100) # 特殊标记
else:
label_name = label[word_idx]["label"]
label_id = label2id[label_name] # 将标签名称转换为 ID
label_ids.append(label_id)
labels.append(label_ids)
tokenized_inputs["labels"] = labels
return tokenized_inputs
# 标签映射
label2id = {
"TIME": 0,
"SCHOOL": 1,
"COLLEGE": 2,
"MAJOR": 3,
"ORG": 4
}
id2label = {v: k for k, v in label2id.items()}
# 加载数据
train_data = load_data("data/train.json")
val_data = load_data("data/val.json")
# 转换为 Hugging Face Dataset
train_dataset = Dataset.from_dict({
"text": [item["text"] for item in train_data],
"labels": [item["labels"] for item in train_data]
})
val_dataset = Dataset.from_dict({
"text": [item["text"] for item in val_data],
"labels": [item["labels"] for item in val_data]
})
# 加载预训练模型和分词器
model_name = "bert-base-chinese"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForTokenClassification.from_pretrained(
model_name, num_labels=len(label2id), id2label=id2label, label2id=label2id
)
# 数据预处理
tokenized_train_dataset = train_dataset.map(
tokenize_and_align_labels, fn_kwargs={"tokenizer": tokenizer}, batched=True
)
tokenized_val_dataset = val_dataset.map(
tokenize_and_align_labels, fn_kwargs={"tokenizer": tokenizer}, batched=True
)
# 定义训练参数
training_args = TrainingArguments(
output_dir="./models/resume_ner_model",
evaluation_strategy="epoch",
learning_rate=2e-5,
per_device_train_batch_size=16,
num_train_epochs=3,
weight_decay=0.01,
save_strategy="epoch",
save_total_limit=2,
logging_dir="./logs",
logging_steps=10,
)
# 定义 Trainer
data_collator = DataCollatorForTokenClassification(tokenizer)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=tokenized_train_dataset,
eval_dataset=tokenized_val_dataset,
tokenizer=tokenizer,
data_collator=data_collator,
)
# 开始训练
trainer.train()
# 保存模型
trainer.save_model("./models/resume_ner_model")
tokenizer.save_pretrained("./models/resume_ner_model")