|
from datasets import load_dataset |
|
from transformers import TrainingArguments |
|
|
|
from span_marker import SpanMarkerModel, Trainer |
|
|
|
|
|
def main() -> None: |
|
|
|
dataset = "Babelscape/multinerd" |
|
train_dataset = load_dataset(dataset, split="train") |
|
eval_dataset = load_dataset(dataset, split="validation").shuffle().select(range(3000)) |
|
labels = [ |
|
"O", |
|
"B-PER", |
|
"I-PER", |
|
"B-ORG", |
|
"I-ORG", |
|
"B-LOC", |
|
"I-LOC", |
|
"B-ANIM", |
|
"I-ANIM", |
|
"B-BIO", |
|
"I-BIO", |
|
"B-CEL", |
|
"I-CEL", |
|
"B-DIS", |
|
"I-DIS", |
|
"B-EVE", |
|
"I-EVE", |
|
"B-FOOD", |
|
"I-FOOD", |
|
"B-INST", |
|
"I-INST", |
|
"B-MEDIA", |
|
"I-MEDIA", |
|
"B-MYTH", |
|
"I-MYTH", |
|
"B-PLANT", |
|
"I-PLANT", |
|
"B-TIME", |
|
"I-TIME", |
|
"B-VEHI", |
|
"I-VEHI", |
|
] |
|
|
|
|
|
model_name = "bert-base-multilingual-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="models/span_marker_mbert_base_multinerd", |
|
|
|
learning_rate=5e-5, |
|
per_device_train_batch_size=32, |
|
per_device_eval_batch_size=32, |
|
|
|
num_train_epochs=1, |
|
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=1000, |
|
save_total_limit=2, |
|
dataloader_num_workers=2, |
|
) |
|
|
|
|
|
trainer = Trainer( |
|
model=model, |
|
args=args, |
|
train_dataset=train_dataset, |
|
eval_dataset=eval_dataset, |
|
) |
|
trainer.train() |
|
trainer.save_model("models/span_marker_mbert_base_multinerd/checkpoint-final") |
|
|
|
test_dataset = load_dataset(dataset, split="test") |
|
|
|
metrics = trainer.evaluate(test_dataset, metric_key_prefix="test") |
|
trainer.save_metrics("test", metrics) |
|
|
|
trainer.create_model_card(language="multilingual", license="apache-2.0") |
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|
|
""" |
|
Logs: |
|
|
|
Training set: |
|
This SpanMarker model will ignore 0.793782% of all annotated entities in the train dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words and the maximum model input length of 256 tokens. |
|
These are the frequencies of the missed entities due to maximum entity length out of 4111958 total entities: |
|
- 12680 missed entities with 9 words (0.308369%) |
|
- 7308 missed entities with 10 words (0.177726%) |
|
- 4414 missed entities with 11 words (0.107345%) |
|
- 2474 missed entities with 12 words (0.060166%) |
|
- 1894 missed entities with 13 words (0.046061%) |
|
- 1130 missed entities with 14 words (0.027481%) |
|
- 744 missed entities with 15 words (0.018094%) |
|
- 582 missed entities with 16 words (0.014154%) |
|
- 344 missed entities with 17 words (0.008366%) |
|
- 226 missed entities with 18 words (0.005496%) |
|
- 84 missed entities with 19 words (0.002043%) |
|
- 46 missed entities with 20 words (0.001119%) |
|
- 20 missed entities with 21 words (0.000486%) |
|
- 20 missed entities with 22 words (0.000486%) |
|
- 12 missed entities with 23 words (0.000292%) |
|
- 18 missed entities with 24 words (0.000438%) |
|
- 2 missed entities with 25 words (0.000049%) |
|
- 4 missed entities with 26 words (0.000097%) |
|
- 4 missed entities with 27 words (0.000097%) |
|
- 2 missed entities with 31 words (0.000049%) |
|
- 8 missed entities with 32 words (0.000195%) |
|
- 6 missed entities with 33 words (0.000146%) |
|
- 2 missed entities with 34 words (0.000049%) |
|
- 4 missed entities with 36 words (0.000097%) |
|
- 8 missed entities with 37 words (0.000195%) |
|
- 2 missed entities with 38 words (0.000049%) |
|
- 2 missed entities with 41 words (0.000049%) |
|
- 2 missed entities with 72 words (0.000049%) |
|
Additionally, a total of 598 (0.014543%) entities were missed due to the maximum input length. |
|
|
|
Validation set: |
|
This SpanMarker model won't be able to predict 0.656224% of all annotated entities in the evaluation dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words. |
|
These are the frequencies of the missed entities due to maximum entity length out of 4724 total entities: |
|
- 19 missed entities with 9 words (0.402202%) |
|
- 7 missed entities with 10 words (0.148180%) |
|
- 1 missed entities with 11 words (0.021169%) |
|
- 1 missed entities with 12 words (0.021169%) |
|
- 2 missed entities with 13 words (0.042337%) |
|
- 1 missed entities with 16 words (0.021169%) |
|
|
|
Testing set: |
|
This SpanMarker model won't be able to predict 0.794755% of all annotated entities in the evaluation dataset. This is caused by the SpanMarkerModel maximum entity length of 8 words and the maximum model input length of 256 tokens. |
|
These are the frequencies of the missed entities due to maximum entity length out of 511856 total entities: |
|
- 1520 missed entities with 9 words (0.296959%) |
|
- 870 missed entities with 10 words (0.169970%) |
|
- 640 missed entities with 11 words (0.125035%) |
|
- 346 missed entities with 12 words (0.067597%) |
|
- 224 missed entities with 13 words (0.043762%) |
|
- 172 missed entities with 14 words (0.033603%) |
|
- 72 missed entities with 15 words (0.014066%) |
|
- 66 missed entities with 16 words (0.012894%) |
|
- 16 missed entities with 17 words (0.003126%) |
|
- 14 missed entities with 18 words (0.002735%) |
|
- 14 missed entities with 19 words (0.002735%) |
|
- 2 missed entities with 20 words (0.000391%) |
|
- 12 missed entities with 21 words (0.002344%) |
|
- 2 missed entities with 24 words (0.000391%) |
|
- 2 missed entities with 25 words (0.000391%) |
|
Additionally, a total of 96 (0.018755%) entities were missed due to the maximum input length. |
|
""" |