|
from datasets import load_dataset |
|
from transformers import TrainingArguments |
|
from span_marker import SpanMarkerModel, Trainer |
|
|
|
|
|
def main() -> None: |
|
|
|
train_dataset = load_dataset("P3ps/Cross_ner", split="train") |
|
test_dataset = load_dataset("P3ps/Cross_ner", split="test") |
|
labels = train_dataset.features["ner_tags"].feature.names |
|
|
|
|
|
model_name = "bert-base-cased" |
|
model = SpanMarkerModel.from_pretrained( |
|
model_name, |
|
labels=labels, |
|
|
|
model_max_length=256, |
|
marker_max_length=128, |
|
entity_max_length=8, |
|
) |
|
|
|
|
|
args = TrainingArguments( |
|
output_dir=f"models/span_marker_bert_base_cross_ner", |
|
run_name=f"bb_cross_ner", |
|
|
|
learning_rate=5e-5, |
|
per_device_train_batch_size=32, |
|
per_device_eval_batch_size=32, |
|
num_train_epochs=3, |
|
weight_decay=0.01, |
|
warmup_ratio=0.1, |
|
bf16=True, |
|
|
|
logging_first_step=True, |
|
logging_steps=50, |
|
evaluation_strategy="steps", |
|
save_strategy="steps", |
|
eval_steps=200, |
|
save_total_limit=2, |
|
dataloader_num_workers=2, |
|
) |
|
|
|
|
|
trainer = Trainer( |
|
model=model, |
|
args=args, |
|
train_dataset=train_dataset, |
|
eval_dataset=test_dataset, |
|
) |
|
trainer.train() |
|
trainer.save_model(f"models/span_marker_bert_base_cross_ner/checkpoint-final") |
|
|
|
|
|
metrics = trainer.evaluate(test_dataset, metric_key_prefix="test") |
|
trainer.save_metrics("test", metrics) |
|
trainer.create_model_card() |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|