File size: 6,327 Bytes
d98c605 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 | """Standalone inference for Process-FM checkpoints (fm-v6 and scaling/sc*).
Dependency: torch + numpy only. The models use global activity/resource ids:
activities 4..776 (UNK=1), resources 2..2739 (UNK=1), log ids from vocab["log_ids"].
"""
import json
from datetime import datetime
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn
N_NUM = 6
ACT_OFFSET = 4
RES_OFFSET = 2
UNK_ACT = 1
UNK_RES = 1
class EventTransformer(nn.Module):
def __init__(self, n_act: int, n_res: int, n_logs: int, n_classes: int,
d_model: int = 256, n_layers: int = 6, n_heads: int = 8,
dropout: float = 0.1, max_len: int = 256,
act_dim: int = 64, res_dim: int = 32, log_dim: int = 16):
super().__init__()
self.act_emb = nn.Embedding(n_act, act_dim, padding_idx=0)
self.res_emb = nn.Embedding(n_res, res_dim, padding_idx=0)
self.log_emb = nn.Embedding(n_logs, log_dim)
self.proj = nn.Linear(act_dim + res_dim + log_dim + N_NUM, d_model)
self.pos = nn.Embedding(max_len, d_model)
layer = nn.TransformerEncoderLayer(
d_model=d_model, nhead=n_heads, dim_feedforward=4 * d_model,
dropout=dropout, batch_first=True, norm_first=True, activation="gelu",
)
self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
self.norm = nn.LayerNorm(d_model)
self.act_out = nn.Linear(d_model, n_classes)
self.res_out = nn.Linear(d_model, n_res)
self.max_len = max_len
def forward(self, gact: torch.Tensor, gres: torch.Tensor, num: torch.Tensor,
log_id: torch.Tensor, causal: bool = True,
return_hidden: bool = False):
b, t = gact.shape
x = torch.cat([
self.act_emb(gact),
self.res_emb(gres),
self.log_emb(log_id).unsqueeze(1).expand(b, t, -1),
num,
], dim=-1)
x = self.proj(x)
idx = torch.arange(t, device=gact.device).unsqueeze(0).expand(b, t)
x = x + self.pos(idx)
if causal:
x = self.encoder(x, mask=torch.triu(
torch.ones(t, t, device=gact.device, dtype=torch.bool), diagonal=1))
else:
x = self.encoder(x)
h = self.norm(x)
if return_hidden:
return h
return self.act_out(h), self.res_out(h)
def build_model(config: dict) -> EventTransformer:
return EventTransformer(
n_act=config["n_act"], n_res=config["n_res"], n_logs=config["n_logs"],
n_classes=config["n_classes"], d_model=config["d_model"],
n_layers=config["n_layers"], n_heads=config["n_heads"],
dropout=config["dropout"], max_len=config["max_len"],
act_dim=config["act_dim"], res_dim=config["res_dim"], log_dim=config["log_dim"],
)
def _to_ns(ts) -> int:
if isinstance(ts, datetime):
return int(ts.timestamp() * 1e9)
return int(np.datetime64(ts, "ns").astype("int64"))
def featurize(events: list[dict], vocab: dict, amount: float = 0.0,
log: str | None = None):
max_len = min(int(vocab.get("max_prefix_len", 256)), 256)
events = events[-max_len:]
n = len(events)
ts = np.array([_to_ns(e["timestamp"]) for e in events], dtype=np.int64)
dt_prev = np.zeros(n, dtype=np.float64)
dt_prev[1:] = (ts[1:] - ts[:-1]) / 1e9
dt_start = (ts - ts[0]) / 1e9
hours = ((ts // 3_600_000_000_000) % 24).astype(np.float64)
weekday = ((ts // 86_400_000_000_000) + 3) % 7
amount_feat = np.log1p(max(amount, 0.0)) / 14.0
num = np.stack([
np.log1p(dt_prev) / 20.0,
np.log1p(dt_start) / 20.0,
np.sin(2 * np.pi * hours / 24.0),
np.cos(2 * np.pi * hours / 24.0),
weekday / 6.0,
np.full(n, amount_feat),
], axis=1).astype(np.float32)
act_ids = {str(a): ACT_OFFSET + i for i, a in enumerate(vocab["activities"])}
res_ids = {str(r): RES_OFFSET + i for i, r in enumerate(vocab["resources"])}
gact = np.array([act_ids.get(str(e.get("activity", "")), UNK_ACT) for e in events],
dtype=np.int64)
gres = np.array([res_ids.get(str(e.get("resource", "")), UNK_RES) for e in events],
dtype=np.int64)
log_name = log or vocab.get("default_log")
log_id = int(vocab["log_ids"].get(log_name, vocab.get("log_unk_id", vocab["n_logs"] - 1)))
return (
torch.from_numpy(gact)[None],
torch.from_numpy(gres)[None],
torch.from_numpy(num)[None],
torch.tensor([log_id], dtype=torch.long),
)
def load_model(release_dir: str | Path = "."):
release_dir = Path(release_dir)
config = json.loads((release_dir / "config.json").read_text())
vocab = json.loads((release_dir / "vocab.json").read_text())
model = build_model(config)
state = torch.load(release_dir / "model.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()
return model, vocab
@torch.no_grad()
def predict(model: EventTransformer, vocab: dict, events: list[dict],
amount: float = 0.0, top_k: int = 5, log: str | None = None) -> list[tuple[str, float]]:
gact, gres, num, log_id = featurize(events, vocab, amount=amount, log=log)
h = model(gact, gres, num, log_id, causal=True, return_hidden=True)
logits = model.act_out(h[0, gact.shape[1] - 1].float())
probs = torch.softmax(logits, dim=-1)
k = min(top_k, probs.numel() - 1)
values, indices = torch.topk(probs, k)
id_to_act = {ACT_OFFSET + i: str(a) for i, a in enumerate(vocab["activities"])}
return [(id_to_act.get(int(i), "UNK"), float(v)) for v, i in zip(values, indices)]
if __name__ == "__main__":
here = Path(__file__).resolve().parent
model, vocab = load_model(here / "fm-v6")
activities = [str(a) for a in vocab["activities"][:3]]
resources = [str(r) for r in vocab["resources"][:2]]
events = [
{"activity": activities[0], "resource": resources[0], "timestamp": "2024-01-01 10:00:00"},
{"activity": activities[1], "resource": resources[1], "timestamp": "2024-01-01 11:00:00"},
{"activity": activities[2], "resource": resources[0], "timestamp": "2024-01-01 12:00:00"},
]
print(predict(model, vocab, events, top_k=3, log=vocab.get("default_log")))
|