flow / modeling_flow.py

Upload modeling_flow.py with huggingface_hub

f34fc9e verified about 1 month ago

2.33 kB

"""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()

	"""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()