Views
No views yet
Note: Full model card with training details coming soon.
model.pt and tokenizer/ from this repo, then:1import torch
2from transformers import AutoModel, AutoTokenizer
3from torch import nn
4
5ENTITY_MARKERS = ["[E1]", "[/E1]", "[E2]", "[/E2]"]
6
7class EventRelationRoBERTa(nn.Module):
8 def __init__(self, model_name, n_new_tokens):
9 super().__init__()
10 self.backbone = AutoModel.from_pretrained(model_name)
11 # Required: training resized the vocab for the 4 entity markers.
12 # Without this, load_state_dict fails on an embedding size mismatch.
13 self.backbone.resize_token_embeddings(self.backbone.config.vocab_size + n_new_tokens)
14 hidden = self.backbone.config.hidden_size
15 self.temporal_head = nn.Linear(hidden, 1)
16 self.causal_head = nn.Linear(hidden, 1)
17
18 def forward(self, input_ids, attention_mask):
19 cls = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
20 return self.temporal_head(cls).squeeze(-1), self.causal_head(cls).squeeze(-1)
21
22tokenizer = AutoTokenizer.from_pretrained("tokenizer/")
23model = EventRelationRoBERTa("roberta-base", len(ENTITY_MARKERS))
24model.load_state_dict(torch.load("model.pt", map_location="cpu", weights_only=True))
25model.eval()1def insert_markers(text, span1, span2):
2 """span1/span2 are [char_start, char_end, ...]; span1 must precede span2."""
3 insertions = sorted([
4 (span1[0], "[E1]"), (span1[1], "[/E1]"),
5 (span2[0], "[E2]"), (span2[1], "[/E2]"),
6 ], key=lambda x: -x[0])
7 for pos, marker in insertions:
8 text = text[:pos] + marker + text[pos:]
9 return text
10
11text = "She opened the door and the cat escaped."
12marked = insert_markers(text, [4, 10], [32, 39])
13# 'She [E1]opened[/E1] the door and the cat [E2]escaped[/E2].'
14
15enc = tokenizer(marked, max_length=256, padding="max_length",
16 truncation=True, return_tensors="pt")
17with torch.no_grad():
18 t_logit, c_logit = model(enc["input_ids"], enc["attention_mask"])
19
20# Raw logits; threshold at 0 (equivalently, sigmoid > 0.5).
21is_sequential = bool(t_logit > 0) # temporal: sequential vs not
22is_causal = bool(c_logit > 0) # causal: causally related vs notmax_length=256: markers pushed past that limit are truncated away and the
prediction becomes meaningless. Check that both markers survive tokenization for
long inputs.1{
2 "model_name": "roberta-base",
3 "max_len": 256,
4 "dims": [
5 "temporal_sequential",
6 "causal"
7 ],
8 "data_source": "/projects/tejo9855/Projects/llm-narrative-annotations/event_relation/outputs/google_gemma-4-31B-it/20260518_143249",
9 "n_train": 6219,
10 "n_val": 690,
11 "val_frac": 0.1,
12 "best_epoch": 4,
13 "seed": 42,
14 "test_f1_gold": 0.805
15}