Upload modeling_flow.py with huggingface_hub

modeling_flow.py

@@ -0,0 +1,72 @@
+"""Transformers-compatible wrapper for the sandbox-lifecycle temporal model.
+
+A GRU over ordered sandbox events, not a transformer. Subclasses
+PreTrainedModel so weights serialize as config.json plus model.safetensors and
+load through AutoModel. Importing this module registers the architecture.
+"""
+
+import torch
+import torch.nn as nn
+from transformers import AutoConfig, AutoModel, PretrainedConfig, PreTrainedModel
+
+EVENTS = [
+    "sandbox_create",
+    "execute",
+    "process_exit",
+    "recover",
+    "wait",
+    "sandbox_teardown",
+]
+LABELS = ["recover_then_wait", "wait_during_outage", "delayed_fault", "cooldown"]
+
+
+class FlowConfig(PretrainedConfig):
+    model_type = "flow_gru"
+
+    def __init__(self, num_labels=4, dropout=0.25, **kwargs):
+        super().__init__(**kwargs)
+        self.num_labels = num_labels
+        self.dropout = dropout
+
+
+class FlowGRU(PreTrainedModel):
+    config_class = FlowConfig
+    base_model_prefix = "flow"
+
+    # Transformers 5 reads this mapping while finalising a load. Nothing here
+    # is tied, but the attribute has to exist or from_pretrained raises
+    # AttributeError before the weights are placed.
+    all_tied_weights_keys = {}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.emb = nn.Embedding(len(EVENTS) + 1, 8, padding_idx=0)
+        self.gru = nn.GRU(input_size=8 + 2, hidden_size=32, batch_first=True)
+        self.head = nn.Linear(32, config.num_labels)
+        self.drop = nn.Dropout(config.dropout)
+
+    def forward(self, events):
+        f = torch.cat([self.emb(events[..., 0].long()), events[..., 1:]], -1)
+        _, h = self.gru(f)
+        return self.head(self.drop(h.squeeze(0)))
+
+    def predict_label(self, sequence):
+        """sequence: list of (type, t, duration) triples. Returns the label."""
+        self.eval()
+        tv = {t: i + 1 for i, t in enumerate(EVENTS)}
+        import numpy as np
+
+        X = np.array(
+            [[float(tv[e[0]]), float(e[1]), float(e[2])] for e in sequence],
+            dtype=np.float32,
+        )
+        with torch.no_grad():
+            return LABELS[int(self.forward(torch.from_numpy(X).unsqueeze(0)).argmax(1)[0])]
+
+
+def register():
+    AutoConfig.register(FlowConfig.model_type, FlowConfig, exist_ok=True)
+    AutoModel.register(FlowConfig, FlowGRU, exist_ok=True)
+
+
+register()