Fine-tuning for Causality Identification¶

Causality Identification is a relation classification task: given a sentence with two entity spans already marked, predict whether a causal relation holds between them.

The dataset column schema for this task is:

Column	Type	Description
text	string	Sentence with entity markers, e.g. <e1>storm</e1> caused <e2>flooding</e2>
relations	list[dict]	Each dict has relationship (0 = no-rel, 1 = causal), first, second

Entity markers (<e1>, </e1>, <e2>, </e2>) are already embedded in text. We register them as special tokens so RoBERTa treats each marker as a single unit rather than splitting it into subwords.

We fine-tune roberta-base as a binary sequence classifier using the HuggingFace Trainer.

Setup¶

In [ ]:

%pip install -q causalatee[huggingface] evaluate

%pip install -q causalatee[huggingface] evaluate

Load the dataset¶

In [ ]:

from datasets import load_dataset

dataset = load_dataset("thagen/AltLex", "causality identification")
print(dataset)

from datasets import load_dataset dataset = load_dataset("thagen/AltLex", "causality identification") print(dataset)

Expected output:

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

Each row's text already contains entity markers:

text      = "The <e1>storm</e1> caused <e2>flooding</e2> in the valley."
relations = [{"relationship": 1, "first": "e1", "second": "e2"}]

Flatten labels¶

Each row may contain multiple entity-pair relations. We derive a single binary label: 1 if any of the listed relations is causal, 0 otherwise.

In [ ]:

def flatten(example):
    is_causal = any(r["relationship"] == 1 for r in example["relations"])
    return {"label": int(is_causal)}


dataset = dataset.map(flatten, remove_columns=["relations"])

def flatten(example): is_causal = any(r["relationship"] == 1 for r in example["relations"]) return {"label": int(is_causal)} dataset = dataset.map(flatten, remove_columns=["relations"])

Tokenize¶

We add the four entity-marker tokens as special tokens so the tokenizer keeps each marker as a single unit (rather than splitting <e1> into <, e, 1, >). The model's embedding table is then resized to accommodate the new tokens.

In [ ]:

from transformers import AutoTokenizer

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

ENTITY_MARKERS = ["<e1>", "</e1>", "<e2>", "</e2>"]
tokenizer.add_special_tokens({"additional_special_tokens": ENTITY_MARKERS})


def tokenize(batch):
    return tokenizer(batch["text"], truncation=True)


tokenized = dataset.map(tokenize, batched=True, remove_columns=["text"])
tokenized.set_format("torch")

from transformers import AutoTokenizer MODEL = "roberta-base" tokenizer = AutoTokenizer.from_pretrained(MODEL) ENTITY_MARKERS = ["<e1>", "</e1>", "<e2>", "</e2>"] tokenizer.add_special_tokens({"additional_special_tokens": ENTITY_MARKERS}) def tokenize(batch): return tokenizer(batch["text"], truncation=True) tokenized = dataset.map(tokenize, batched=True, remove_columns=["text"]) tokenized.set_format("torch")

Define the model¶

After adding special tokens, resize_token_embeddings initializes new embedding rows for the four entity markers (the existing weights are preserved).

In [ ]:

from transformers import AutoModelForSequenceClassification

id2label = {0: "no-rel", 1: "causal"}
label2id = {v: k for k, v in id2label.items()}

model = AutoModelForSequenceClassification.from_pretrained(
    MODEL,
    num_labels=2,
    id2label=id2label,
    label2id=label2id,
)
model.resize_token_embeddings(len(tokenizer))

from transformers import AutoModelForSequenceClassification id2label = {0: "no-rel", 1: "causal"} label2id = {v: k for k, v in id2label.items()} model = AutoModelForSequenceClassification.from_pretrained( MODEL, num_labels=2, id2label=id2label, label2id=label2id, ) model.resize_token_embeddings(len(tokenizer))

Evaluation metric¶

In [ ]:

import evaluate
import numpy as np

f1_metric = evaluate.load("f1")


def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return f1_metric.compute(predictions=predictions, references=labels, average="binary")

import evaluate import numpy as np f1_metric = evaluate.load("f1") def compute_metrics(eval_pred): logits, labels = eval_pred predictions = np.argmax(logits, axis=-1) return f1_metric.compute(predictions=predictions, references=labels, average="binary")

Train¶

In [ ]:

from transformers import DataCollatorWithPadding, Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="roberta-causality-identification",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="f1",
    num_train_epochs=3,
    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=DataCollatorWithPadding(tokenizer),
    compute_metrics=compute_metrics,
)

trainer.train()

from transformers import DataCollatorWithPadding, Trainer, TrainingArguments training_args = TrainingArguments( output_dir="roberta-causality-identification", eval_strategy="epoch", save_strategy="epoch", load_best_model_at_end=True, metric_for_best_model="f1", num_train_epochs=3, 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=DataCollatorWithPadding(tokenizer), compute_metrics=compute_metrics, ) trainer.train()

Expected training output:

{'eval_loss': 0.4569, 'eval_f1': 0.9829, 'epoch': 1.0}
{'loss': 0.5597, 'grad_norm': 1.246, 'learning_rate': 1.96e-05, 'epoch': 1.316}
{'eval_loss': 0.006028, 'eval_f1': 0.9957, 'epoch': 2.0}
{'loss': 0.02074, 'grad_norm': 0.03467, 'learning_rate': 4.688e-06, 'epoch': 2.632}
{'eval_loss': 0.0021, 'eval_f1': 1.0, 'epoch': 3.0}

Note: AltLex is a dataset of causal sentences by construction. After flattening to a binary causal / no-rel label, virtually all examples are positive, so F1 converges to 1.0 quickly. For a more challenging benchmark with balanced classes, use a dataset that includes non-causal pairs such as SemEval2010T8.

Evaluate¶

In [ ]:

results = trainer.evaluate()
print(results)

results = trainer.evaluate() print(results)

Expected output:

{'eval_loss': 0.002100, 'eval_f1': 1.0, 'eval_runtime': 0.192,
 'eval_samples_per_second': 2100.4, 'eval_steps_per_second': 67.6, 'epoch': 3.0}

Use the model¶

The trained model is compatible with causalatee's CausalityIdentificationPipeline.

In [ ]:

from transformers import pipeline

from causalatee.integrations.huggingface import CausalityIdentificationPipeline  # noqa: F401 -- registers the pipeline

pipe = pipeline(
    "causality-identification",
    model=trainer.model,
    tokenizer=tokenizer,
)
pipe("The <e1>heavy rainfall</e1> caused <e2>widespread flooding</e2> across the valley.")

from transformers import pipeline from causalatee.integrations.huggingface import CausalityIdentificationPipeline # noqa: F401 -- registers the pipeline pipe = pipeline( "causality-identification", model=trainer.model, tokenizer=tokenizer, ) pipe("The <e1>heavy rainfall</e1> caused <e2>widespread flooding</e2> across the valley.")

Expected output:

{'relationship': 'causal', 'score': 0.961}