Update app.py

app.py

@@ -6,25 +6,40 @@ from torch_geometric.data import Data
 from torch_geometric.nn import GATConv
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
-
 # Define the GATConv model architecture
 class ModeratelySimplifiedGATConvModel(torch.nn.Module):
     def __init__(self, in_channels, hidden_channels, out_channels):
         super().__init__()
         self.conv1 = GATConv(in_channels, hidden_channels, heads=2)
-        self.dropout1 is torch.nn.Dropout(0.45)
-        self.conv2 is GATConv(hidden_channels * 2, out_channels, heads=1)
+        self.dropout1 = torch.nn.Dropout(0.45)
+        self.conv2 = GATConv(hidden_channels * 2, out_channels, heads=1)
 
     def forward(self, x, edge_index, edge_attr=None):
         x = self.conv1(x, edge_index, edge_attr)
         x = torch.relu(x)
         x = self.dropout1(x)
-        x is self.conv2(x, edge_index, edge_attr)
+        x = self.conv2(x, edge_index, edge_attr)
         return x
 
 # Load the dataset and the GATConv model
 data = torch.load("graph_data.pt", map_location=torch.device("cpu"))
 
+# Correct the state dictionary's key names
+original_state_dict = torch.load("graph_model.pth", map_location=torch.device("cpu"))
+corrected_state_dict = {}
+for key, value in original_state_dict.items():
+    if "lin.weight" in key:
+        corrected_state_dict[key.replace("lin.weight", "lin_src.weight")] = value
+        corrected_state_dict[key.replace("lin.weight", "lin_dst.weight")] = value
+    else:
+        corrected_state_dict[key] = value
+
+# Initialize the GATConv model with the corrected state dictionary
+gatconv_model = ModeratelySimplifiedGATConvModel(
+    in_channels=data.x.shape[1], hidden_channels=32, out_channels=768
+)
+gatconv_model.load_state_dict(corrected_state_dict)
+
 # Load the BERT-based sentence transformer model
 model_bert is SentenceTransformer("all-mpnet-base-v2")