| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from torch_geometric.data import Data |
| from torch_geometric.nn import GATv2Conv, LayerNorm |
| from safetensors.torch import load_file |
|
|
| class BipartiteData(Data): |
| """bipartite graph data structure for pyg batching compatibility""" |
| def __init__(self, x_var=None, x_con=None, edge_index_c2v=None, edge_attr=None, **kwargs): |
| super().__init__(**kwargs) |
| self.x_var = x_var |
| self.x_con = x_con |
| self.edge_index_c2v = edge_index_c2v |
| self.edge_attr = edge_attr |
| if x_var is not None and x_con is not None: |
| self.num_nodes = x_var.size(0) + x_con.size(0) |
|
|
| def __inc__(self, key, value, *args, **kwargs): |
| if key == 'edge_index_c2v': |
| return torch.tensor([[self.x_con.size(0)], [self.x_var.size(0)]]) |
| return super().__inc__(key, value, *args, **kwargs) |
|
|
| class BipartiteTransformerLayer(nn.Module): |
| """alternating bipartite graph attention layer""" |
| def __init__(self, hidden_dim, heads=8, dropout=0.1): |
| super().__init__() |
| self.attn_c2v = GATv2Conv((hidden_dim, hidden_dim), hidden_dim // heads, heads=heads, edge_dim=1, dropout=dropout, add_self_loops=False) |
| self.norm_v = LayerNorm(hidden_dim) |
| self.ffn_v = nn.Sequential(nn.Linear(hidden_dim, hidden_dim * 2), nn.GELU(), nn.Linear(hidden_dim * 2, hidden_dim)) |
| self.norm_ffn_v = LayerNorm(hidden_dim) |
| |
| self.attn_v2c = GATv2Conv((hidden_dim, hidden_dim), hidden_dim // heads, heads=heads, edge_dim=1, dropout=dropout, add_self_loops=False) |
| self.norm_c = LayerNorm(hidden_dim) |
| self.ffn_c = nn.Sequential(nn.Linear(hidden_dim, hidden_dim * 2), nn.GELU(), nn.Linear(hidden_dim * 2, hidden_dim)) |
| self.norm_ffn_c = LayerNorm(hidden_dim) |
|
|
| def forward(self, x_var, x_con, edge_index_c2v, edge_attr): |
| h_var = self.attn_c2v((x_con, x_var), edge_index_c2v, edge_attr=edge_attr) |
| x_var = self.norm_v(x_var + F.dropout(h_var, p=0.1, training=self.training)) |
| x_var = self.norm_ffn_v(x_var + self.ffn_v(x_var)) |
|
|
| edge_index_v2c = edge_index_c2v.flip(0) |
| h_con = self.attn_v2c((x_var, x_con), edge_index_v2c, edge_attr=edge_attr) |
| x_con = self.norm_c(x_con + F.dropout(h_con, p=0.1, training=self.training)) |
| x_con = self.norm_ffn_c(x_con + self.ffn_c(x_con)) |
|
|
| return x_var, x_con |
|
|
| class LippersheyBase(nn.Module): |
| """lippershey bipartite graph transformer base model""" |
| def __init__(self, node_in_dim=5, hidden_dim=256, num_layers=6, heads=8): |
| super().__init__() |
| self.var_proj = nn.Linear(node_in_dim, hidden_dim) |
| self.con_proj = nn.Linear(node_in_dim, hidden_dim) |
| self.layers = nn.ModuleList([BipartiteTransformerLayer(hidden_dim, heads) for _ in range(num_layers)]) |
| self.mlm_head = nn.Sequential(nn.Linear(hidden_dim, hidden_dim), nn.GELU(), nn.Linear(hidden_dim, 1)) |
|
|
| def forward(self, x_var, x_con, edge_index_c2v, edge_attr): |
| v = self.var_proj(x_var) |
| c = self.con_proj(x_con) |
| for layer in self.layers: |
| v, c = layer(v, c, edge_index_c2v, edge_attr) |
| return self.mlm_head(v) |
|
|
| def extract_features(self, x_var, x_con, edge_index_c2v, edge_attr): |
| """extract latent representations of variables for downstream ppo policy""" |
| self.eval() |
| with torch.no_grad(): |
| v = self.var_proj(x_var) |
| c = self.con_proj(x_con) |
| for layer in self.layers: |
| v, c = layer(v, c, edge_index_c2v, edge_attr) |
| return v |
|
|
| @classmethod |
| def from_pretrained_safetensors(cls, safetensors_path, node_in_dim=5, hidden_dim=256, num_layers=6, heads=8): |
| """load model directly from safetensors format""" |
| model = cls(node_in_dim=node_in_dim, hidden_dim=hidden_dim, num_layers=num_layers, heads=heads) |
| state_dict = load_file(safetensors_path) |
| model.load_state_dict(state_dict) |
| return model |
