from datasets import load_dataset from transformers import TrainingArguments from span_marker import SpanMarkerModel, Trainer def main() -> None: # Load the dataset, ensure "tokens" and "ner_tags" columns, and get a list of labels 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", ] # Initialize a SpanMarker model using a pretrained BERT-style encoder model_name = "bert-base-multilingual-cased" model = SpanMarkerModel.from_pretrained( model_name, labels=labels, # SpanMarker hyperparameters: model_max_length=256, marker_max_length=128, entity_max_length=8, ) # Prepare the 🤗 transformers training arguments args = TrainingArguments( output_dir="models/span_marker_mbert_base_multinerd", # Training Hyperparameters: learning_rate=5e-5, per_device_train_batch_size=32, per_device_eval_batch_size=32, # gradient_accumulation_steps=2, num_train_epochs=1, weight_decay=0.01, warmup_ratio=0.1, bf16=True, # Replace `bf16` with `fp16` if your hardware can't use bf16. # Other Training parameters logging_first_step=True, logging_steps=50, evaluation_strategy="steps", save_strategy="steps", eval_steps=1000, save_total_limit=2, dataloader_num_workers=2, ) # Initialize the trainer using our model, training args & dataset, and train 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") # Compute & save the metrics on the test set 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. """