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 = "xlm-roberta-base" model = SpanMarkerModel.from_pretrained( model_name, labels=labels, # SpanMarker hyperparameters: model_max_length=256, marker_max_length=128, entity_max_length=6, ) # Prepare the 🤗 transformers training arguments args = TrainingArguments( output_dir="models/span_marker_xlm_roberta_base_multinerd", # Training Hyperparameters: learning_rate=1e-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_xlm_roberta_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) if __name__ == "__main__": main() """ This SpanMarker model will ignore 2.239322% of all annotated entities in the train dataset. This is caused by the SpanMarkerModel maximum entity length of 6 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: - 35814 missed entities with 7 words (0.870972%) - 21246 missed entities with 8 words (0.516688%) - 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 2978 (0.072423%) entities were missed due to the maximum input length. This SpanMarker model won't be able to predict 2.501087% of all annotated entities in the evaluation dataset. This is caused by the SpanMarkerModel maximum entity length of 6 words. These are the frequencies of the missed entities due to maximum entity length out of 4598 total entities: - 45 missed entities with 7 words (0.978686%) - 27 missed entities with 8 words (0.587212%) - 21 missed entities with 9 words (0.456720%) - 9 missed entities with 10 words (0.195737%) - 3 missed entities with 12 words (0.065246%) - 4 missed entities with 13 words (0.086994%) - 3 missed entities with 14 words (0.065246%) - 1 missed entities with 15 words (0.021749%) - 1 missed entities with 16 words (0.021749%) - 1 missed entities with 20 words (0.021749%) """ """ wandb: Run summary: wandb: eval/loss 0.00594 wandb: eval/overall_accuracy 0.98181 wandb: eval/overall_f1 0.90333 wandb: eval/overall_precision 0.91259 wandb: eval/overall_recall 0.89427 wandb: eval/runtime 21.4308 wandb: eval/samples_per_second 154.171 wandb: eval/steps_per_second 4.853 wandb: test/loss 0.00559 wandb: test/overall_accuracy 0.98247 wandb: test/overall_f1 0.91314 wandb: test/overall_precision 0.91994 wandb: test/overall_recall 0.90643 wandb: test/runtime 2202.6894 wandb: test/samples_per_second 169.652 wandb: test/steps_per_second 5.302 wandb: train/epoch 1.0 wandb: train/global_step 93223 wandb: train/learning_rate 0.0 wandb: train/loss 0.0049 wandb: train/total_flos 7.851073325660897e+17 wandb: train/train_loss 0.01782 wandb: train/train_runtime 41756.9748 wandb: train/train_samples_per_second 71.44 wandb: train/train_steps_per_second 2.233 """