SpanMarker

The repository provides several high-performance pretrained models for Named Entity Recognition (NER) across different datasets. These models can be used via the Hugging Face Inference API or loaded locally.

Available Model Families:

• FewNERD: High-performance models for the fine-grained Few-NERD dataset, including bert-base, roberta-large, and multilingual xlm-roberta-base versions.
• OntoNotes v5.0: A roberta-large based model (tomaarsen/span-marker-roberta-large-ontonotes5) that outperforms standard spaCy transformer models.
• CoNLL03: Models using xlm-roberta-large, including versions optimized with document-level context for higher F1 scores.
• CoNLL++: A competitive model for the CoNLL++ dataset using xlm-roberta-large and document-level context.
• MultiNERD: Multilingual models trained on the MultiNERD dataset.
  • tomaarsen/span-marker-xlm-roberta-base-multinerd (requires separating punctuation from words for best performance).
  • tomaarsen/span-marker-mbert-base-multinerd (successor, uses bert-base-multilingual-cased, no punctuation limitation).

You can integrate any SpanMarker model from the Hugging Face Hub into a spaCy pipeline. To do this, load a standard spaCy model (excluding the default ner component) and add the span_marker pipeline component using nlp.add_pipe(). You must provide a config dictionary containing the specific Hugging Face model name.

import spacy

# Load the spaCy model with the span_marker pipeline component
nlp = spacy.load("en_core_web_sm", exclude=["ner"])
nlp.add_pipe("span_marker", config={"model": "tomaarsen/span-marker-roberta-large-ontonotes5"})

# Feed some text through the model to get a spacy Doc
text = """Cleopatra VII, also known as Cleopatra the Great, was the last active ruler of the \" 
"""Ptolemaic Kingdom of Egypt. She was born in 69 BCE and ruled Egypt from 51 BCE until her \" 
"""death in 30 BCE."""
doc = nlp(text)

# And look at the entities
print([(entity, entity.label_) for entity in doc.ents])

You can train a Named Entity Recognition (NER) model using SpanMarker by following these steps:

1. Load your dataset: Ensure your dataset contains the necessary columns for training.
2. Define labels: Provide a list of entity labels. SpanMarker automatically handles various tagging schemes like IOB, IOB2, BIOES, or BILOU, or even datasets with no specific scheme.
3. Initialize SpanMarkerModel: Use SpanMarkerModel.from_pretrained() with a pretrained encoder (e.g., bert-base-cased) and your list of labels.
4. Configure TrainingArguments: Use standard Hugging Face TrainingArguments to define hyperparameters like learning rate, batch size, and epochs.
5. Use the Trainer: Initialize the SpanMarker Trainer (which subclasses the Hugging Face Trainer) with your model, arguments, and datasets, then call .train().
6. Evaluate and Save: Use .evaluate() to check performance and .save_model() or .push_to_hub() to persist your results.

from datasets import load_dataset
from span_marker import SpanMarkerModel, Trainer
from transformers import TrainingArguments

# 1. Load dataset
dataset = load_dataset("DFKI-SLT/few-nerd", "supervised")
labels = ["O", "art", "building", "event", "location", "organization", "other", "person", "product"]

# 2. Initialize model with an encoder and labels
encoder_id = "bert-base-cased"
model = SpanMarkerModel.from_pretrained(encoder_id, labels=labels)

# 3. Set up training arguments
args = TrainingArguments(
    output_dir="my_span_marker_model",
    learning_rate=5e-5,
    gradient_accumulation_steps=2,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=1,
    save_strategy="steps",
    eval_steps=200,
    logging_steps=50,
    warmup_ratio=0.1,
)

# 4. Initialize and run the Trainer
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=dataset["train"].select(range(8000)),
    eval_dataset=dataset["validation"].select(range(2000)),
)

trainer.train()

# 5. Evaluate and Save
metrics = trainer.evaluate()
print(metrics)

trainer.save_model("my_span_marker_model/checkpoint-final")
trainer.push_to_hub("my_span_marker_model/checkpoint-final")

Training a SpanMarker model involves initializing a SpanMarkerModel with a pretrained encoder and labels, configuring transformers.TrainingArguments, and using the span_marker.Trainer.

Key Steps:

1. Prepare Dataset: Ensure your dataset has tokens and ner_tags columns. Supported annotation schemes include IOB, IOB2, BIOES, and BILOU.
2. Initialize Model: Use SpanMarkerModel.from_pretrained(encoder_id, labels=labels, ...).
3. Configure Hyperparameters:
   • model_max_length: Maximum length for the encoder.
   • marker_max_length: Maximum length for the marker.
   • entity_max_length: Maximum length for an entity span.
   • model_card_data: Use SpanMarkerModelCardData to provide metadata for the Hugging Face Hub.
4. Train: Use the Trainer class from span_marker to execute the training loop.

from pathlib import Path
from datasets import load_dataset
from transformers import TrainingArguments
from span_marker import SpanMarkerModel, Trainer, SpanMarkerModelCardData

def main() -> None:
    # 1. Load and prepare dataset
    dataset_id = "DFKI-SLT/few-nerd"
    dataset_name = "FewNERD"
    dataset = load_dataset(dataset_id, "supervised")
    dataset = dataset.remove_columns("ner_tags")
    dataset = dataset.rename_column("fine_ner_tags", "ner_tags")
    labels = dataset["train"].features["ner_tags"].feature.names

    # 2. Initialize SpanMarker model
    encoder_id = "bert-base-cased"
    model_id = f"tomaarsen/span-marker-{encoder_id}-fewnerd-fine-super"
    model = SpanMarkerModel.from_pretrained(
        encoder_id,
        labels=labels,
        model_max_length=256,
        marker_max_length=128,
        entity_max_length=8,
        model_card_data=SpanMarkerModelCardData(
            model_id=model_id,
            encoder_id=encoder_id,
            dataset_name=dataset_name,
            dataset_id=dataset_id,
            license="cc-by-sa-4.0",
            language="en",
        ),
    )

    # 3. Prepare training arguments
    output_dir = Path("models") / model_id
    args = TrainingArguments(
        output_dir=output_dir,
        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,  # Use fp16 if bf16 is not supported
        logging_first_step=True,
        logging_steps=50,
        eval_strategy="steps",
        save_strategy="steps",
        eval_steps=3000,
        save_total_limit=2,
        dataloader_num_workers=2,
    )

    # 4. Initialize and run trainer
    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["validation"],
    )
    trainer.train()

    # 5. Evaluate and save
    metrics = trainer.evaluate(dataset["test"], metric_key_prefix="test")
    trainer.save_metrics("test", metrics)
    trainer.save_model(output_dir / "checkpoint-final")

if __name__ == "__main__":
    main()

SpanMarker is a Named Entity Recognition (NER) module that fine-tunes pretrained encoders (like BERT or RoBERTa) using a marker-based approach. Instead of traditional token-level classification, it uses special start and end tokens to represent spans.

Core Workflow:

1. Tokenization: The input text is tokenized and padded.
2. Marker Insertion: Special tokens <start> (ID 50261) and <end> (ID 50262) are appended to represent every legal span in the sentence.
3. Positioning: Position IDs are updated to virtually place these markers between the text tokens.
4. Attention Masking: A custom attention mask matrix is used to allow markers to attend to text tokens (one-directional attention) while preventing text tokens from attending to markers.
5. Embedding & Classification: The encoder produces embeddings for all tokens. For each span, the model concatenates the embedding of its <start> marker and its <end> marker. This concatenated vector is passed through a linear layer to predict the entity label using cross-entropy loss.