Skip to content

Commit a8c703b

Browse files
kmontemayor2-sckmontemayorclaudekmonte
authored
Make example scripts self-consistent in naming and distributed behavior (#688)
Co-authored-by: kmontemayor <kmontemayor@snapchat.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Kyle Montemayor <kyle.e.montemayor@gmail.com>
1 parent a69defe commit a8c703b

16 files changed

Lines changed: 469 additions & 605 deletions

‎examples/link_prediction/cora.ipynb‎

Lines changed: 1 addition & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -399,48 +399,7 @@
399399
"id": "54158f1c",
400400
"metadata": {},
401401
"outputs": [],
402-
"source": [
403-
"### Examining the trained model\n",
404-
"# The trained model will be saved at: `GbmlConfig.SharedConfig.TrainedModelMetadata.trained_model_uri`\n",
405-
"print(\n",
406-
" frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n",
407-
")\n",
408-
"\n",
409-
"# You can load the model locally and play around with it:\n",
410-
"import torch\n",
411-
"from torch_geometric.data import Data\n",
412-
"\n",
413-
"from examples.link_prediction.models import init_example_gigl_homogeneous_model\n",
414-
"from gigl.common import UriFactory\n",
415-
"from gigl.src.common.utils.model import load_state_dict_from_uri\n",
416-
"\n",
417-
"node_type = frozen_config.graph_metadata_pb_wrapper.homogeneous_condensed_node_type\n",
418-
"edge_type = frozen_config.graph_metadata_pb_wrapper.homogeneous_condensed_edge_type\n",
419-
"node_feature_dim = frozen_config.preprocessed_metadata_pb_wrapper.condensed_node_type_to_feature_dim_map[\n",
420-
" node_type\n",
421-
"]\n",
422-
"edge_feature_dim = frozen_config.preprocessed_metadata_pb_wrapper.condensed_edge_type_to_feature_dim_map[\n",
423-
" edge_type\n",
424-
"]\n",
425-
"model = init_example_gigl_homogeneous_model(\n",
426-
" node_feature_dim=node_feature_dim,\n",
427-
" edge_feature_dim=edge_feature_dim,\n",
428-
" device=torch.device(\"cpu\"),\n",
429-
" state_dict=load_state_dict_from_uri(\n",
430-
" UriFactory.create_uri(\n",
431-
" frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n",
432-
" )\n",
433-
" ),\n",
434-
")\n",
435-
"print(model)\n",
436-
"\n",
437-
"# Create some random data to test the model.\n",
438-
"example_data = Data(\n",
439-
" x=torch.rand((10, node_feature_dim)), edge_index=torch.randint(0, 10, (2, 20))\n",
440-
")\n",
441-
"embeddings = model(example_data, device=torch.device(\"cpu\"))\n",
442-
"print(embeddings)"
443-
]
402+
"source": "### Examining the trained model\n# The trained model will be saved at: `GbmlConfig.SharedConfig.TrainedModelMetadata.trained_model_uri`\nprint(\n frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n)\n\n# You can load the model locally and play around with it:\nimport torch\nfrom torch_geometric.data import Data\n\nfrom examples.link_prediction.models import init_example_gigl_homogeneous_model\nfrom gigl.common import UriFactory\nfrom gigl.src.common.utils.model import load_state_dict_from_uri\n\nnode_feature_dim = frozen_config.node_type_to_feature_dim_map[\n frozen_config.graph_metadata_pb_wrapper.homogeneous_node_type\n]\nedge_feature_dim = frozen_config.edge_type_to_feature_dim_map[\n frozen_config.graph_metadata_pb_wrapper.homogeneous_edge_type\n]\nmodel = init_example_gigl_homogeneous_model(\n node_feature_dim=node_feature_dim,\n edge_feature_dim=edge_feature_dim,\n device=torch.device(\"cpu\"),\n state_dict=load_state_dict_from_uri(\n UriFactory.create_uri(\n frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n )\n ),\n)\nprint(model)\n\n# Create some random data to test the model.\nexample_data = Data(\n x=torch.rand((10, node_feature_dim)), edge_index=torch.randint(0, 10, (2, 20))\n)\nembeddings = model(example_data, device=torch.device(\"cpu\"))\nprint(embeddings)"
444403
},
445404
{
446405
"cell_type": "markdown",

‎examples/link_prediction/dblp.ipynb‎

Lines changed: 1 addition & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -396,70 +396,7 @@
396396
"id": "54158f1c",
397397
"metadata": {},
398398
"outputs": [],
399-
"source": [
400-
"### Examining the trained model\n",
401-
"# The trained model will be saved at: `GbmlConfig.SharedConfig.TrainedModelMetadata.trained_model_uri`\n",
402-
"print(\n",
403-
" frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n",
404-
")\n",
405-
"\n",
406-
"# You can load the model locally and play around with it:\n",
407-
"import torch\n",
408-
"from torch_geometric.data import HeteroData\n",
409-
"\n",
410-
"from examples.link_prediction.models import init_example_gigl_heterogeneous_model\n",
411-
"from gigl.common import UriFactory\n",
412-
"from gigl.src.common.types.graph_data import EdgeType, NodeType\n",
413-
"from gigl.src.common.utils.model import load_state_dict_from_uri\n",
414-
"\n",
415-
"edge_types = frozen_config.task_metadata_pb_wrapper.get_supervision_edge_types()\n",
416-
"print(f\"Supervision edge types: {edge_types}\")\n",
417-
"edge_type = edge_types[0]\n",
418-
"query_node_type = edge_type.src_node_type\n",
419-
"labeled_node_type = edge_type.dst_node_type\n",
420-
"\n",
421-
"graph_metadata = frozen_config.graph_metadata_pb_wrapper\n",
422-
"\n",
423-
"# Build dicts of a specific node type and edge type to feature dimension.\n",
424-
"# This dimension is the *input* dimension of the model.\n",
425-
"node_feature_dims: dict[NodeType, int] = {\n",
426-
" graph_metadata.condensed_node_type_to_node_type_map[\n",
427-
" condensed_node_type\n",
428-
" ]: node_feature_dim\n",
429-
" for condensed_node_type, node_feature_dim in frozen_config.preprocessed_metadata_pb_wrapper.condensed_node_type_to_feature_dim_map.items()\n",
430-
"}\n",
431-
"edge_feature_dims: dict[EdgeType, int] = {\n",
432-
" graph_metadata.condensed_edge_type_to_edge_type_map[\n",
433-
" condensed_edge_type\n",
434-
" ]: edge_feature_dim\n",
435-
" for condensed_edge_type, edge_feature_dim in frozen_config.preprocessed_metadata_pb_wrapper.condensed_edge_type_to_feature_dim_map.items()\n",
436-
"}\n",
437-
"model = init_example_gigl_heterogeneous_model(\n",
438-
" node_type_to_feature_dim=node_feature_dims,\n",
439-
" edge_type_to_feature_dim=edge_feature_dims,\n",
440-
" device=torch.device(\"cpu\"),\n",
441-
" state_dict=load_state_dict_from_uri(\n",
442-
" UriFactory.create_uri(\n",
443-
" frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n",
444-
" )\n",
445-
" ),\n",
446-
")\n",
447-
"print(model)\n",
448-
"\n",
449-
"# Create some random data to test the model.\n",
450-
"example_data = HeteroData()\n",
451-
"example_data[query_node_type].x = torch.rand((10, node_feature_dims[node_type]))\n",
452-
"example_data[labeled_node_type].x = torch.rand(\n",
453-
" (10, node_feature_dims[labeled_node_type])\n",
454-
")\n",
455-
"example_data[edge_type].edge_index = torch.randint(0, 10, (2, 20))\n",
456-
"embeddings = model(\n",
457-
" example_data,\n",
458-
" device=torch.device(\"cpu\"),\n",
459-
" output_node_types=[edge_type.src_node_type, edge_type.dst_node_type],\n",
460-
")\n",
461-
"print(embeddings)"
462-
]
399+
"source": "### Examining the trained model\n# The trained model will be saved at: `GbmlConfig.SharedConfig.TrainedModelMetadata.trained_model_uri`\nprint(\n frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n)\n\n# You can load the model locally and play around with it:\nimport torch\nfrom torch_geometric.data import HeteroData\n\nfrom examples.link_prediction.models import init_example_gigl_heterogeneous_model\nfrom gigl.common import UriFactory\nfrom gigl.src.common.utils.model import load_state_dict_from_uri\n\nedge_types = frozen_config.task_metadata_pb_wrapper.get_supervision_edge_types()\nprint(f\"Supervision edge types: {edge_types}\")\nedge_type = edge_types[0]\nquery_node_type = edge_type.src_node_type\nlabeled_node_type = edge_type.dst_node_type\n\ngraph_metadata = frozen_config.graph_metadata_pb_wrapper\n\n# Map each node/edge type to its feature dimension, which is the *input*\n# dimension of the model.\nnode_feature_dims = frozen_config.node_type_to_feature_dim_map\nedge_feature_dims = frozen_config.edge_type_to_feature_dim_map\nmodel = init_example_gigl_heterogeneous_model(\n node_type_to_feature_dim=node_feature_dims,\n edge_type_to_feature_dim=edge_feature_dims,\n device=torch.device(\"cpu\"),\n state_dict=load_state_dict_from_uri(\n UriFactory.create_uri(\n frozen_config.gbml_config_pb.shared_config.trained_model_metadata.trained_model_uri\n )\n ),\n)\nprint(model)\n\n# Create some random data to test the model.\nexample_data = HeteroData()\nexample_data[query_node_type].x = torch.rand((10, node_feature_dims[node_type]))\nexample_data[labeled_node_type].x = torch.rand(\n (10, node_feature_dims[labeled_node_type])\n)\nexample_data[edge_type].edge_index = torch.randint(0, 10, (2, 20))\nembeddings = model(\n example_data,\n device=torch.device(\"cpu\"),\n output_node_types=[edge_type.src_node_type, edge_type.dst_node_type],\n)\nprint(embeddings)"
463400
},
464401
{
465402
"cell_type": "markdown",

0 commit comments

Comments
 (0)