Update pages/19_Graphs3.py

pages/19_Graphs3.py

@@ -155,10 +155,10 @@ def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
         return node_set["empty_state"]
     raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
 
-def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
+# Define the GNN model function
+def gnn_model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
-    graph = tfgnn.keras.layers.MapFeatures(
-        node_sets_fn=set_initial_node_states)(graph)
+    graph = tfgnn.keras.layers.MapFeatures(node_sets_fn=set_initial_node_states)(graph)
     for i in range(num_graph_updates):
         graph = mt_albis.MtAlbisGraphUpdate(
             units=node_state_dim,
@@ -171,10 +171,17 @@ def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
             normalization_type="layer",
             next_state_type="residual",
         )(graph)
-    graph = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="feat")(graph)
-    graph = tf.keras.layers.Dense(349, activation="softmax")(graph)
     return tf.keras.Model(inputs, graph)
 
+# Define the complete model function
+def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
+    gnn_model = gnn_model_fn(graph_tensor_spec)
+    inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
+    graph = gnn_model(inputs)
+    paper_state = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="hidden_state")(graph)
+    paper_state = tf.keras.layers.Dense(349, activation="softmax")(paper_state)
+    return tf.keras.Model(inputs, paper_state)
+
 # Check for TPU/ GPU and set strategy
 st.write("Setting up strategy for distributed training...")
 if tf.config.list_physical_devices("TPU"):
@@ -227,48 +234,11 @@ trainer = runner.KerasTrainer(
     backup_and_restore=False,
 )
 
-# Define feature processors
-st.write("Defining feature processors...")
-def process_node_features(node_set: tfgnn.NodeSet, node_set_name: str):
-    if node_set_name == "field_of_study":
-        return {"hashed_id": tf.keras.layers.Hashing(50_000)(node_set["#id"])}
-    if node_set_name == "institution":
-        return {"hashed_id": tf.keras.layers.Hashing(6_500)(node_set["#id"])}
-    if node_set_name == "paper":
-        return {"feat": node_set["feat"], "label": node_set["label"]}
-    if node_set_name == "author":
-        return {"empty_state": tfgnn.keras.layers.MakeEmptyFeature()(node_set)}
-    raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
-
-def drop_all_features(_, **unused_kwargs):
-    return {}
-
-process_features = tfgnn.keras.layers.MapFeatures(
-    context_fn=drop_all_features,
-    node_sets_fn=process_node_features,
-    edge_sets_fn=drop_all_features)
-
-add_readout = tfgnn.keras.layers.AddReadoutFromFirstNode("seed", node_set_name="paper")
-move_label_to_readout = tfgnn.keras.layers.StructuredReadoutIntoFeature(
-    "seed", feature_name="label", new_feature_name="paper_venue", remove_input_feature=True)
-
-feature_processors = [process_features, add_readout, move_label_to_readout]
-
-# Function to create the full model with feature processing
-def create_full_model(graph_tensor_spec: tfgnn.GraphTensorSpec):
-    input_graph = tf.keras.layers.Input(type_spec=graph_tensor_spec)
-    graph = input_graph
-    for processor in feature_processors:
-        graph = processor(graph)
-    graph = tfgnn.keras.layers.Readout(node_set_name="paper", feature_name="feat")(graph)
-    output_graph = model_fn(graph_tensor_spec)(graph)
-    return tf.keras.Model(input_graph, output_graph)
-
 # Run training
 st.write("Training the model...")
 runner.run(
     task=task,
-    model_fn=create_full_model,
+    model_fn=model_fn,
     trainer=trainer,
     optimizer_fn=optimizer_fn,
     epochs=epochs,