Update app.py

app.py

@@ -6,7 +6,6 @@ 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):
@@ -18,13 +17,29 @@ class ModeratelySimplifiedGATConvModel(torch.nn.Module):
     def forward(self, x, edge_index, edge_attr=None):
         x = self.conv1(x, edge_index, edge_attr)
         x = torch.relu(x)
-        x = dropout1(x)
+        x = self.dropout1(x)
         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 = SentenceTransformer("all-mpnet-base-v2")