Fine-tuning for Causal Candidate Extraction¶

Causal Candidate Extraction is a span extraction task: given a sentence, identify all character spans that correspond to causal entities (causes and effects).

The dataset column schema for this task is:

Column	Type	Description
text	string	Input sentence
entity	list[list[int]]	List of [char_start, char_end] spans

We fine-tune roberta-base as a token classifier using BIO labels. The spans in entity are character-level, so we need a preprocessing step to convert them to per-token BIO labels aligned with the tokenizer's subword offsets.

Evaluation uses causalatee's span-overlap metrics (precision, recall, F1, and granularity-penalised F1) — matching the Touché shared-task evaluator.

Setup¶

In [ ]:

%pip install -q causalatee[huggingface]

%pip install -q causalatee[huggingface]

Load the dataset¶

In [ ]:

from datasets import load_dataset

dataset = load_dataset("thagen/AltLex", "causal candidate extraction")
print(dataset)

from datasets import load_dataset dataset = load_dataset("thagen/AltLex", "causal candidate extraction") print(dataset)

Expected output:

DatasetDict({
    train: Dataset({features: ['text', 'entity'], num_rows: 1984})
    test:  Dataset({features: ['text', 'entity'], num_rows: 496})
})

Each row's entity field is a list of [char_start, char_end] pairs, for example:

text   = "The storm caused flooding in the valley."
entity = [[4, 9], [16, 24]]  # "storm" and "flooding"

Convert character spans to BIO token labels¶

RoBERTa operates on subword tokens, not characters. We tokenize each sentence with return_offsets_mapping=True to get the character range (char_start, char_end) for every subword, then assign:

• B-SPAN — first token of a gold span
• I-SPAN — continuation token inside a gold span
• O — outside any span
• -100 — special tokens (CLS, SEP, PAD) are ignored by the loss

We also capture the eval split's offset mappings before removing that column, because span_compute_metrics needs them at evaluation time to decode BIO predictions back into character spans.

In [ ]:

from transformers import AutoTokenizer

MODEL = "roberta-base"
tokenizer = AutoTokenizer.from_pretrained(MODEL)

LABEL2ID = {"O": 0, "B-SPAN": 1, "I-SPAN": 2}
ID2LABEL = {v: k for k, v in LABEL2ID.items()}


def char_spans_to_bio(char_start, char_end, gold_spans):
    """Return BIO label id for a single token given its character range."""
    if char_start == char_end:  # special token
        return -100
    for span_start, span_end in gold_spans:
        if char_start >= span_start and char_end <= span_end:
            return LABEL2ID["B-SPAN"] if char_start == span_start else LABEL2ID["I-SPAN"]
    return LABEL2ID["O"]


def tokenize_and_align(batch):
    encoding = tokenizer(
        batch["text"],
        truncation=True,
        return_offsets_mapping=True,
    )
    all_labels = []
    for offsets, gold_spans in zip(encoding["offset_mapping"], batch["entity"]):
        labels = [char_spans_to_bio(cs, ce, gold_spans) for cs, ce in offsets]
        all_labels.append(labels)
    encoding["labels"] = all_labels
    return encoding


# Capture eval offsets before the column is removed
eval_offset_mappings = dataset["test"].map(
    lambda batch: tokenizer(batch["text"], truncation=True, return_offsets_mapping=True),
    batched=True,
)["offset_mapping"]

tokenized = dataset.map(
    tokenize_and_align,
    batched=True,
    remove_columns=dataset["train"].column_names,
)
tokenized.set_format("torch")

from transformers import AutoTokenizer MODEL = "roberta-base" tokenizer = AutoTokenizer.from_pretrained(MODEL) LABEL2ID = {"O": 0, "B-SPAN": 1, "I-SPAN": 2} ID2LABEL = {v: k for k, v in LABEL2ID.items()} def char_spans_to_bio(char_start, char_end, gold_spans): """Return BIO label id for a single token given its character range.""" if char_start == char_end: # special token return -100 for span_start, span_end in gold_spans: if char_start >= span_start and char_end <= span_end: return LABEL2ID["B-SPAN"] if char_start == span_start else LABEL2ID["I-SPAN"] return LABEL2ID["O"] def tokenize_and_align(batch): encoding = tokenizer( batch["text"], truncation=True, return_offsets_mapping=True, ) all_labels = [] for offsets, gold_spans in zip(encoding["offset_mapping"], batch["entity"]): labels = [char_spans_to_bio(cs, ce, gold_spans) for cs, ce in offsets] all_labels.append(labels) encoding["labels"] = all_labels return encoding # Capture eval offsets before the column is removed eval_offset_mappings = dataset["test"].map( lambda batch: tokenizer(batch["text"], truncation=True, return_offsets_mapping=True), batched=True, )["offset_mapping"] tokenized = dataset.map( tokenize_and_align, batched=True, remove_columns=dataset["train"].column_names, ) tokenized.set_format("torch")

Define the model¶

In [ ]:

from transformers import AutoModelForTokenClassification

model = AutoModelForTokenClassification.from_pretrained(
    MODEL,
    num_labels=len(ID2LABEL),
    id2label=ID2LABEL,
    label2id=LABEL2ID,
)

from transformers import AutoModelForTokenClassification model = AutoModelForTokenClassification.from_pretrained( MODEL, num_labels=len(ID2LABEL), id2label=ID2LABEL, label2id=LABEL2ID, )

Evaluation metric¶

span_compute_metrics decodes BIO predictions back to character-level spans and computes macro-averaged precision, recall, F1, and granularity-penalised F1 — the same metrics used by the Touché shared-task evaluator. The granularity penalty follows Potthast et al. (2013), discounting F1 when predictions are over-fragmented relative to gold spans.

In [ ]:

from causalatee.integrations.huggingface import span_compute_metrics

compute_metrics = span_compute_metrics(model.config.id2label, eval_offset_mappings)

from causalatee.integrations.huggingface import span_compute_metrics compute_metrics = span_compute_metrics(model.config.id2label, eval_offset_mappings)

Train¶

In [ ]:

from transformers import DataCollatorForTokenClassification, Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="roberta-causal-candidate-extraction",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    num_train_epochs=5,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=32,
    learning_rate=2e-5,
    warmup_steps=50,
    fp16=True,
    logging_steps=50,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized["train"],
    eval_dataset=tokenized["test"],
    data_collator=DataCollatorForTokenClassification(tokenizer),
    compute_metrics=compute_metrics,
)

trainer.train()

from transformers import DataCollatorForTokenClassification, Trainer, TrainingArguments training_args = TrainingArguments( output_dir="roberta-causal-candidate-extraction", eval_strategy="epoch", save_strategy="epoch", load_best_model_at_end=True, metric_for_best_model="f1", num_train_epochs=5, per_device_train_batch_size=16, per_device_eval_batch_size=32, learning_rate=2e-5, warmup_steps=50, fp16=True, logging_steps=50, ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized["train"], eval_dataset=tokenized["test"], data_collator=DataCollatorForTokenClassification(tokenizer), compute_metrics=compute_metrics, ) trainer.train()

Expected training output (values will vary):

{'loss': 0.4312, 'learning_rate': 1.8e-05, 'epoch': 1.0}
{'eval_loss': 0.2941, 'eval_precision': 0.712, 'eval_recall': 0.681, 'eval_f1': 0.696, 'eval_f1_gran': 0.672, 'epoch': 1.0}
{'loss': 0.1853, 'learning_rate': 6e-06, 'epoch': 3.0}
{'eval_loss': 0.2204, 'eval_precision': 0.783, 'eval_recall': 0.761, 'eval_f1': 0.772, 'eval_f1_gran': 0.748, 'epoch': 3.0}
{'loss': 0.0914, 'learning_rate': 0.0, 'epoch': 5.0}
{'eval_loss': 0.2398, 'eval_precision': 0.801, 'eval_recall': 0.778, 'eval_f1': 0.789, 'eval_f1_gran': 0.762, 'epoch': 5.0}

Evaluate¶

In [ ]:

results = trainer.evaluate()
print(results)

results = trainer.evaluate() print(results)

Expected output:

{'eval_loss': 0.2398, 'eval_precision': 0.801, 'eval_recall': 0.778,
 'eval_f1': 0.789, 'eval_granularity': 1.031, 'eval_f1_gran': 0.762,
 'eval_intersection_over_union': 0.714, 'eval_runtime': 5.43}

Use the model¶

The trained model is compatible with causalatee's CausalCandidateExtractionPipeline.

In [ ]:

from transformers import pipeline

from causalatee.integrations.huggingface import (
    CausalCandidateExtractionPipeline,  # noqa: F401 -- registers the pipeline
)

pipe = pipeline(
    "causal-candidate-extraction",
    model=trainer.model,
    tokenizer=tokenizer,
)
pipe("The heavy rainfall caused widespread flooding across the valley.")

from transformers import pipeline from causalatee.integrations.huggingface import ( CausalCandidateExtractionPipeline, # noqa: F401 -- registers the pipeline ) pipe = pipeline( "causal-candidate-extraction", model=trainer.model, tokenizer=tokenizer, ) pipe("The heavy rainfall caused widespread flooding across the valley.")

Expected output:

[{'start': 4, 'end': 19, 'entity': 'SPAN'},
 {'start': 27, 'end': 44, 'entity': 'SPAN'}]