narrative-event-relation-roberta

RoBERTa-base fine-tuned for event relation identification: binary temporal (sequential vs non-sequential) and causal classification. Uses entity markers [E1]/[E2] inserted at character offsets. Trained on LLM pseudo-labels with held-out human gold evaluation.

Note: Full model card with training details coming soon.

Loading

Download model.pt and tokenizer/ from this repo, then:

import torch
from transformers import AutoModel, AutoTokenizer
from torch import nn

ENTITY_MARKERS = ["[E1]", "[/E1]", "[E2]", "[/E2]"]

class EventRelationRoBERTa(nn.Module):
    def __init__(self, model_name, n_new_tokens):
        super().__init__()
        self.backbone = AutoModel.from_pretrained(model_name)
        # Required: training resized the vocab for the 4 entity markers.
        # Without this, load_state_dict fails on an embedding size mismatch.
        self.backbone.resize_token_embeddings(self.backbone.config.vocab_size + n_new_tokens)
        hidden = self.backbone.config.hidden_size
        self.temporal_head = nn.Linear(hidden, 1)
        self.causal_head   = nn.Linear(hidden, 1)

    def forward(self, input_ids, attention_mask):
        cls = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
        return self.temporal_head(cls).squeeze(-1), self.causal_head(cls).squeeze(-1)

tokenizer = AutoTokenizer.from_pretrained("tokenizer/")
model = EventRelationRoBERTa("roberta-base", len(ENTITY_MARKERS))
model.load_state_dict(torch.load("model.pt", map_location="cpu", weights_only=True))
model.eval()

Input format

The model takes a single string with the two candidate event triggers wrapped in entity markers at their character offsets. It never sees spans as structured input.

def insert_markers(text, span1, span2):
    """span1/span2 are [char_start, char_end, ...]; span1 must precede span2."""
    insertions = sorted([
        (span1[0], "[E1]"), (span1[1], "[/E1]"),
        (span2[0], "[E2]"), (span2[1], "[/E2]"),
    ], key=lambda x: -x[0])
    for pos, marker in insertions:
        text = text[:pos] + marker + text[pos:]
    return text

text = "She opened the door and the cat escaped."
marked = insert_markers(text, [4, 10], [32, 39])
# 'She [E1]opened[/E1] the door and the cat [E2]escaped[/E2].'

enc = tokenizer(marked, max_length=256, padding="max_length",
                truncation=True, return_tensors="pt")
with torch.no_grad():
    t_logit, c_logit = model(enc["input_ids"], enc["attention_mask"])

# Raw logits; threshold at 0 (equivalently, sigmoid > 0.5).
is_sequential = bool(t_logit > 0)   # temporal: sequential vs not
is_causal     = bool(c_logit > 0)   # causal: causally related vs not

Note max_length=256: markers pushed past that limit are truncated away and the prediction becomes meaningless. Check that both markers survive tokenization for long inputs.

Config

{
  "model_name": "roberta-base",
  "max_len": 256,
  "dims": [
    "temporal_sequential",
    "causal"
  ],
  "data_source": "/projects/tejo9855/Projects/llm-narrative-annotations/event_relation/outputs/google_gemma-4-31B-it/20260518_143249",
  "n_train": 6219,
  "n_val": 690,
  "val_frac": 0.1,
  "best_epoch": 4,
  "seed": 42,
  "test_f1_gold": 0.805
}