Update pages/15_Graphs.py

pages/15_Graphs.py

@@ -12,39 +12,36 @@ os.environ['TF_USE_LEGACY_KERAS'] = '1'
 # Define the model function
 def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = inputs = tf.keras.Input(type_spec=graph_tensor_spec)
-    
-    # Encode input features
+
+    # Encode input features (callback omitted for brevity).
     graph = tfgnn.keras.layers.MapFeatures(
         node_sets_fn=lambda node_set, node_set_name: tf.keras.layers.Dense(64)(node_set['features'])
     )(graph)
-    
-    # Message passing layers
+
+    # For each round of message passing...
     for _ in range(2):
+        # ... create and apply a Keras layer.
         graph = mt_albis.MtAlbisGraphUpdate(
-            units=128,
-            message_dim=64,
-            receiver_tag=tfgnn.NODES,  # Ensure it is correctly set to tfgnn.NODES
-            attention_type="none",
-            simple_conv_reduce_type="mean",
-            normalization_type="layer",
-            next_state_type="residual",
-            state_dropout_rate=0.2,
-            l2_regularization=1e-5,
+            units=128, message_dim=64,
+            attention_type="none", simple_conv_reduce_type="mean",
+            normalization_type="layer", next_state_type="residual",
+            state_dropout_rate=0.2, l2_regularization=1e-5,
+            receiver_tag=tfgnn.NODES  # Correctly use tfgnn.NODES
         )(graph)
-    
+
     return tf.keras.Model(inputs, graph)
 
 # Function to create a sample graph
 def create_sample_graph():
     num_nodes = 10
     num_edges = 15
-    
+
     graph = nx.gnm_random_graph(num_nodes, num_edges, directed=True)
-    
+
     # Create a GraphTensor
     node_features = tf.random.normal((num_nodes, 64))
     edge_features = tf.random.normal((num_edges, 32))
-    
+
     graph_tensor = tfgnn.GraphTensor.from_pieces(
         node_sets={
             "papers": tfgnn.NodeSet.from_fields(
@@ -63,32 +60,32 @@ def create_sample_graph():
             )
         }
     )
-    
+
     return graph, graph_tensor
 
 # Streamlit app
 def main():
     st.title("Graph Neural Network Architecture Visualization")
-    
+
     # Create sample graph
     nx_graph, graph_tensor = create_sample_graph()
-    
+
     # Create and compile the model
     model = model_fn(graph_tensor.spec)
     model.compile(optimizer='adam', loss='binary_crossentropy')
-    
+
     # Display model summary
     st.subheader("Model Summary")
     model.summary(print_fn=lambda x: st.text(x))
-    
+
     # Visualize the graph
     st.subheader("Sample Graph Visualization")
     fig, ax = plt.subplots(figsize=(10, 8))
     pos = nx.spring_layout(nx_graph)
-    nx.draw(nx_graph, pos, with_labels=True, node_color='lightblue', 
+    nx.draw(nx_graph, pos, with_labels=True, node_color='lightblue',
             node_size=500, arrowsize=20, ax=ax)
     st.pyplot(fig)
-    
+
     # Display graph tensor info
     st.subheader("Graph Tensor Information")
     st.text(f"Number of nodes: {graph_tensor.node_sets['papers'].total_size}")