Fine-tuning for Causality Detection¶

Causality Detection is a binary sentence classification task: given a sentence, predict whether it expresses a causal relation (causal) or not (uncausal).

The dataset column schema for this task is:

Column	Type	Description
text	string	Input sentence
label	ClassLabel {0: uncausal, 1: causal}	Ground-truth label

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

Setup¶

Install the library with its HuggingFace extras and the evaluate library for metrics.

In [ ]:

%pip install -q causalatee[huggingface] evaluate

%pip install -q causalatee[huggingface] evaluate

Load the dataset¶

All causalatee datasets are hosted on the HuggingFace Hub. Each dataset exposes multiple task-specific configurations; pass the config name as the second argument to load_dataset.

In [ ]:

from datasets import load_dataset

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

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

Expected output:

DatasetDict({
    train: Dataset({features: ['index', 'text', 'label'], num_rows: 596})
    test:  Dataset({features: ['index', 'text', 'label'], num_rows: 404})
})

Tokenize¶

RoBERTa uses a byte-level BPE tokenizer. We truncate to the model's maximum length and pad to the longest sequence in each batch (handled automatically by the Trainer).

In [ ]:

from transformers import AutoTokenizer

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


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


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

from transformers import AutoTokenizer MODEL = "roberta-base" tokenizer = AutoTokenizer.from_pretrained(MODEL) def tokenize(batch): return tokenizer(batch["text"], truncation=True) tokenized = dataset.map(tokenize, batched=True, remove_columns=["index", "text"]) tokenized.set_format("torch")

Define the model¶

We load roberta-base as a sequence classifier and attach the label mapping so that the trained model can produce human-readable predictions without extra post-processing.

In [ ]:

from transformers import AutoModelForSequenceClassification

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

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

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

Evaluation metric¶

Causality detection is evaluated with binary F1 (positive class = causal).

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¶

The TrainingArguments below use sensible defaults for a classification fine-tune. Adjust num_train_epochs, learning_rate, and per_device_train_batch_size to match your available hardware.

In [ ]:

from transformers import DataCollatorWithPadding, Trainer, TrainingArguments

training_args = TrainingArguments(
    output_dir="roberta-causality-detection",
    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-detection", 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.6817, 'eval_f1': 0.0, 'epoch': 1.0}
{'loss': 0.6924, 'grad_norm': 2.547, 'learning_rate': 1.96e-05, 'epoch': 1.316}
{'eval_loss': 0.7360, 'eval_f1': 0.4674, 'epoch': 2.0}
{'loss': 0.6405, 'grad_norm': 8.478, 'learning_rate': 4.688e-06, 'epoch': 2.632}
{'eval_loss': 0.6791, 'eval_f1': 0.4664, 'epoch': 3.0}

Note: AltLex has only 596 training examples with a skewed test split (28% causal), so results on this dataset will be modest. Larger datasets like BECauSEv2 or combined training across multiple corpora will yield significantly better F1.

Evaluate¶

After training, run a final evaluation on the test split.

In [ ]:

results = trainer.evaluate()
print(results)

results = trainer.evaluate() print(results)

Expected output:

{'eval_loss': 0.7360, 'eval_f1': 0.4674, 'eval_runtime': 0.186,
 'eval_samples_per_second': 2173.1, 'eval_steps_per_second': 69.9, 'epoch': 3.0}

Use the model¶

The trained model is compatible with causalatee's CausalityDetectionPipeline, which wraps it for convenient inference.

In [ ]:

from transformers import pipeline

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

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

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

Expected output:

{'label': 'causal', 'score': 0.973}