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
5class NarrativeRoBERTa(nn.Module):
6 def __init__(self, model_name, n_dims):
7 super().__init__()
8 self.backbone = AutoModel.from_pretrained(model_name)
9 hidden = self.backbone.config.hidden_size
10 self.heads = nn.ModuleList([nn.Linear(hidden, 1) for _ in range(n_dims)])
11
12 def forward(self, input_ids, attention_mask):
13 cls = self.backbone(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0, :]
14 return torch.cat([h(cls) for h in self.heads], dim=1)
15
16tokenizer = AutoTokenizer.from_pretrained("tokenizer/")
17model = NarrativeRoBERTa("roberta-base", n_dims=9)
18model.load_state_dict(torch.load("model.pt", map_location="cpu", weights_only=True))
19model.eval()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}