Rename pages/19_Graphs.py to pages/19_Graphs3.py

pages/{19_Graphs.py → 19_Graphs3.py}

@@ -6,6 +6,7 @@ from tensorflow_gnn.experimental import sampler
 from tensorflow_gnn.models import mt_albis
 import functools
 import os
+from typing import Mapping
 
 # Set environment variable for legacy Keras
 os.environ['TF_USE_LEGACY_KERAS'] = '1'
@@ -13,9 +14,6 @@ os.environ['TF_USE_LEGACY_KERAS'] = '1'
 # Set Streamlit title
 st.title("Solving OGBN-MAG end-to-end with TF-GNN")
 
-# Install necessary packages
-st.write("Installing necessary packages...")
-
 st.write("Setting up the environment...")
 tf.get_logger().setLevel('ERROR')
 st.write(f"Running TF-GNN {tfgnn.__version__} under TensorFlow {tf.__version__}.")
@@ -140,6 +138,23 @@ example_input_graph_spec = train_ds_provider.input_graph_spec._unbatch()
 st.write("Dataset providers created successfully.")
 
 # Define the model function
+node_state_dim = 128
+num_graph_updates = 4
+message_dim = 128
+state_dropout_rate = 0.2
+l2_regularization = 1e-5
+
+def set_initial_node_states(node_set: tfgnn.NodeSet, node_set_name: str):
+    if node_set_name == "field_of_study":
+        return tf.keras.layers.Embedding(50_000, 32)(node_set["hashed_id"])
+    if node_set_name == "institution":
+        return tf.keras.layers.Embedding(6_500, 16)(node_set["hashed_id"])
+    if node_set_name == "paper":
+        return tf.keras.layers.Dense(node_state_dim, "relu")(node_set["feat"])
+    if node_set_name == "author":
+        return node_set["empty_state"]
+    raise KeyError(f"Unexpected node_set_name='{node_set_name}'")
+
 def model_fn(graph_tensor_spec: tfgnn.GraphTensorSpec):
     graph = inputs = tf.keras.layers.Input(type_spec=graph_tensor_spec)
     graph = tfgnn.keras.layers.MapFeatures(