File size: 1,294 Bytes
45507bb | 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 | """Numeric and log state encoders.
NumericStateEncoder: a fixed-length feature vector (metrics, sensor readings, a
tabular row) -> one token per feature (value * feature embedding + feature id),
so the core can attend to individual features. Missing values are NaN -> masked.
LogStateEncoder: a text state encoder for logs. Logs are text, so it reuses the
text encoder; it only normalises volatile tokens (hex ids, numbers, timestamps)
and keeps the last `max_lines` lines before tokenisation (see minijev.text_task_utils).
"""
from __future__ import annotations
import torch
from torch import nn
class NumericStateEncoder(nn.Module):
uses_backbone = False
def __init__(self, d: int, n_features: int, **_):
super().__init__()
self.feat = nn.Parameter(torch.zeros(n_features, d))
self.scale = nn.Parameter(torch.zeros(n_features, d))
nn.init.normal_(self.feat, std=0.02)
nn.init.normal_(self.scale, std=0.02)
self.mlp = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, d))
def forward(self, x): # (B, F) float, NaN = missing
valid = ~torch.isnan(x)
v = torch.nan_to_num(x).unsqueeze(-1)
tok = self.feat + v * self.scale
tok = tok + self.mlp(tok)
return tok, valid
|