Update logic2.py

logic2.py

@@ -22,9 +22,56 @@ import yaml
 
 import pickle
 #----------------------------------------------
+# SAGE model
+class GNNEncoder(torch.nn.Module):
+    def __init__(self, hidden_channels, out_channels):
+        super().__init__()
+        # these convolutions have been replicated to match the number of edge types
+        self.conv1 = SAGEConv((-1, -1), hidden_channels)
+        self.conv2 = SAGEConv((-1, -1), out_channels)
 
+    def forward(self, x, edge_index):
+        x = self.conv1(x, edge_index).relu()
+        x = self.conv2(x, edge_index)
+        return x
+
+class EdgeDecoder(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.lin1 = Linear(2 * hidden_channels, hidden_channels)
+        self.lin2 = Linear(hidden_channels, 1)
+        
+    def forward(self, z_dict, edge_label_index):
+        row, col = edge_label_index
+        # concat user and movie embeddings
+        z = torch.cat([z_dict['user'][row], z_dict['movie'][col]], dim=-1)
+        # concatenated embeddings passed to linear layer
+        z = self.lin1(z).relu()
+        z = self.lin2(z)
+        return z.view(-1)
+
+class Model(torch.nn.Module):
+    def __init__(self, hidden_channels):
+        super().__init__()
+        self.encoder = GNNEncoder(hidden_channels, hidden_channels)
+        self.encoder = to_hetero(self.encoder, data.metadata(), aggr='sum')
+        self.decoder = EdgeDecoder(hidden_channels)
+
+    def forward(self, x_dict, edge_index_dict, edge_label_index):
+        # z_dict contains dictionary of movie and user embeddings returned from GraphSage
+        z_dict = self.encoder(x_dict, edge_index_dict)
+        return self.decoder(z_dict, edge_label_index)
+#----------------------------------------------
 def load_hetero_data():
     with open('Hgraph.pkl', 'rb') as file:
         data = pickle.load(file)
         return data
+
+def load_model(train_data, val_data, test_data):
+    model = Model(hidden_channels=32)
+    with torch.no_grad():
+        model.encoder(train_data.x_dict, train_data.edge_index_dict)
+    model.load_state_dict(torch.load('model.pt',map_location=torch.device('cpu')))
+    model.eval()
+    return model