Showing preview only (352K chars total). Download the full file or copy to clipboard to get everything.
Repository: google-research/talk-like-a-graph
Branch: main
Commit: 36af51e19ef7
Files: 20
Total size: 340.3 KB
Directory structure:
gitextract_nuchrmnk/
├── CONTRIBUTING.md
├── LICENSE
├── README.md
├── __init__.py
├── talk_like_a_graph/
│ ├── README.md
│ ├── __init__.py
│ ├── graph_generators.py
│ ├── graph_generators_runner.py
│ ├── graph_generators_test.py
│ ├── graph_metrics.py
│ ├── graph_metrics_test.py
│ ├── graph_tasks.py
│ ├── graph_tasks_generator.py
│ ├── graph_tasks_utils.py
│ ├── graph_text_encoders.py
│ ├── graph_text_encoders_test.py
│ └── name_dictionaries.py
└── tutorial/
├── KDD-Tutorial-1-Talk-Like-a-Graph.ipynb
├── KDD-Tutorial-2-Let-Your-Graph-Do-The-Talking.ipynb
└── README.md
================================================
FILE CONTENTS
================================================
================================================
FILE: CONTRIBUTING.md
================================================
# How to Contribute
## Contributor License Agreement
Contributions to this project must be accompanied by a Contributor License
Agreement. You (or your employer) retain the copyright to your contribution,
this simply gives us permission to use and redistribute your contributions as
part of the project. Head over to <https://cla.developers.google.com/> to see
your current agreements on file or to sign a new one.
You generally only need to submit a CLA once, so if you've already submitted one
(even if it was for a different project), you probably don't need to do it
again.
## Code reviews
All submissions, including submissions by project members, require review. We
use GitHub pull requests for this purpose. Consult
[GitHub Help](https://help.github.com/articles/about-pull-requests/) for more
information on using pull requests.
## Community Guidelines
This project follows [Google's Open Source Community
Guidelines](https://opensource.google/conduct/).
================================================
FILE: LICENSE
================================================
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
================================================
FILE: README.md
================================================
# Encoding Graphs and Structured Data in Language Models
This repository contains code for
[Talk like a Graph: Encoding Graphs for Large Language Models](https://arxiv.org/abs/2310.04560)
and
[Let Your Graph Do the Talking: Encoding Structured Data for LLMs](https://arxiv.org/abs/2402.05862).
## Cite us
If you use this package for published work, please cite the following.
For the "Talk like a Graph" project:
```
@inproceedigs{fatemi2024talk,
title={Talk like a Graph: Encoding Graphs for Large Language Models},
author={Bahare Fatemi and Jonathan Halcrow and Bryan Perozzi},
booktitle={International Conference on Learning Representations (ICLR)},
year={2024}
}
```
For Graph Token (a.k.a "Let Your Graph Do the Talking"):
```
@misc{perozzi2024letgraphtalkingencoding,
title={Let Your Graph Do the Talking: Encoding Structured Data for LLMs},
author={Bryan Perozzi and Bahare Fatemi and Dustin Zelle and Anton Tsitsulin and Mehran Kazemi and Rami Al-Rfou and Jonathan Halcrow},
year={2024},
eprint={2402.05862},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2402.05862},
}
```
## Disclaimer
This is not an official Google product.
================================================
FILE: __init__.py
================================================
================================================
FILE: talk_like_a_graph/README.md
================================================
# Using Large Language Models to Solve Graph Problems
This repository contains the code to generate graph reasoning problems with
different graph generator algorithms and graph encoding methods, as well as
different prompting techniques.
The graph tasks are `edge existence`, `node degree`, `node count`, `edge count`,
`connected nodes`, `disconnected nodes`, `cycle check`, `reachability`,
`shortest path`, `maximum flow`, `node classification`, and `triangle counting`.
The datasets used here are proposed in our paper:
[Talk like a Graph: Encoding Graphs for Large Language Models](https://arxiv.org/abs/2310.04560).
### Generating graphs
```sh
./graphqa/graph_generator.sh
```
### Generating files for tasks
```sh
./graphqa/task_generator.sh
```
## Contact us
For questions or comments about the implementation, please contact
baharef@google.com.
## Cite us
If you use this package for published work, please cite the following:
```
@inproceedigs{fatemi2024talk,
title={Talk like a Graph: Encoding Graphs for Large Language Models},
author={Bahare Fatemi and Jonathan Halcrow and Bryan Perozzi},
booktitle={International Conference on Learning Representations (ICLR)},
year={2024}
}
```
## Disclaimer
This is not an official Google product.
# Placeholder for internal data notes.
================================================
FILE: talk_like_a_graph/__init__.py
================================================
================================================
FILE: talk_like_a_graph/graph_generators.py
================================================
r"""Random graph generation."""
import random
import networkx as nx
import numpy as np
_NUMBER_OF_NODES_RANGE = {
"small": np.arange(5, 10),
"medium": np.arange(10, 15),
"large": np.arange(15, 20),
}
_NUMBER_OF_COMMUNITIES_RANGE = {
"small": np.arange(2, 4),
"medium": np.arange(2, 8),
"large": np.arange(2, 10),
}
def generate_graphs(
number_of_graphs: int,
algorithm: str,
directed: bool,
random_seed: int = 1234,
er_min_sparsity: float = 0.0,
er_max_sparsity: float = 1.0,
) -> list[nx.Graph]:
"""Generating multiple graphs using the provided algorithms.
Args:
number_of_graphs: number of graphs to generate
algorithm: the random graph generator algorithm
directed: whether to generate directed or undirected graphs.
random_seed: the random seed to generate graphs with.
er_min_sparsity: minimum sparsity of er graphs.
er_max_sparsity: maximum sparsity of er graphs.
Returns:
generated_graphs: a list of nx graphs.
Raises:
NotImplementedError: if the algorithm is not yet implemented.
"""
random.seed(random_seed)
np.random.seed(random_seed)
generated_graphs = []
graph_sizes = random.choices(
list(_NUMBER_OF_NODES_RANGE.keys()), k=number_of_graphs
)
random_state = np.random.RandomState(random_seed)
if algorithm == "er":
for i in range(number_of_graphs):
sparsity = random.uniform(er_min_sparsity, er_max_sparsity)
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
generated_graphs.append(
nx.erdos_renyi_graph(
number_of_nodes, sparsity, seed=random_state, directed=directed
)
)
elif algorithm == "ba":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
m = random.randint(1, number_of_nodes - 1)
generated_graph = nx.barabasi_albert_graph(
number_of_nodes, m, seed=random_state
)
if directed:
generated_graphs.append(randomize_directions(generated_graph))
else:
generated_graphs.append(generated_graph)
elif algorithm == "sbm":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
number_of_communities = random.choice(
_NUMBER_OF_COMMUNITIES_RANGE[graph_sizes[i]]
)
# sizes forms number of nodes in communities.
sizes = []
for _ in range(number_of_communities - 1):
sizes.append(
random.randint(
1,
max(
1,
number_of_nodes - sum(sizes) - (number_of_communities - 1),
),
)
)
sizes.append(number_of_nodes - sum(sizes))
# p forms probabilities of communities connecting each other.
p = np.random.uniform(size=(number_of_communities, number_of_communities))
if random.uniform(0, 1) < 0.5:
p = np.maximum(p, p.transpose())
else:
p = np.minimum(p, p.transpose())
sbm_graph = nx.stochastic_block_model(
sizes, p, seed=random_state, directed=directed
)
# sbm graph generator automatically adds dictionary attributes.
sbm_graph = remove_graph_data(sbm_graph)
generated_graphs.append(sbm_graph)
elif algorithm == "sfn":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
generated_graph = nx.scale_free_graph(number_of_nodes, seed=random_state)
# sfn graphs are by defaukt directed.
if not directed:
generated_graphs.append(remove_directions(generated_graph))
else:
generated_graphs.append(generated_graph)
elif algorithm == "complete":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
create_using = nx.DiGraph if directed else nx.Graph
generated_graphs.append(
nx.complete_graph(number_of_nodes, create_using=create_using)
)
elif algorithm == "star":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
# number_of_nodes for star is the input + a center node.
generated_graph = nx.star_graph(number_of_nodes - 1)
if directed:
generated_graphs.append(randomize_directions(generated_graph))
else:
generated_graphs.append(generated_graph)
elif algorithm == "path":
for i in range(number_of_graphs):
number_of_nodes = random.choice(_NUMBER_OF_NODES_RANGE[graph_sizes[i]])
create_using = nx.DiGraph if directed else nx.Graph
generated_graphs.append(
nx.path_graph(number_of_nodes, create_using=create_using)
)
else:
raise NotImplementedError()
return generated_graphs
def remove_graph_data(graph: nx.Graph) -> nx.Graph:
# GraphML writer does not support dictionary data for nodes or graphs.
for ind in range((graph.number_of_nodes())):
graph.nodes[ind].pop("block", None)
graph_data_keys = list(graph.graph.keys())
for _, node in enumerate(graph_data_keys):
graph.graph.pop(node, None)
return graph
def randomize_directions(graph: nx.Graph) -> nx.DiGraph:
# Converting the undirected graph to a directed graph.
directed_graph = graph.to_directed()
# For each edge, randomly choose a direction.
edges = list(graph.edges())
for u, v in edges:
if random.random() < 0.5:
directed_graph.remove_edge(u, v)
else:
directed_graph.remove_edge(v, u)
return directed_graph
def remove_directions(graph: nx.Graph) -> nx.Graph:
# Converting the direted graph to an undirected one by removing directions.
undirected_graph = nx.Graph()
undirected_graph.add_nodes_from(graph.nodes())
# Add edges between nodes, ignoring directions.
for u, v in graph.edges():
undirected_graph.add_edge(u, v)
return undirected_graph
================================================
FILE: talk_like_a_graph/graph_generators_runner.py
================================================
r"""Random graph generation.
This code generates random graph using different algorithms.
# Placeholder for Google-internal comments.
"""
from collections.abc import Sequence
import os
from absl import app
from absl import flags
import networkx as nx
# Internal import.
from . import graph_generators
_ALGORITHM = flags.DEFINE_string(
"algorithm",
None,
"The graph generating algorithm to use.",
required=True,
)
_NUMBER_OF_GRAPHS = flags.DEFINE_integer(
"number_of_graphs",
None,
"The number of graphs to generate.",
required=True,
)
_DIRECTED = flags.DEFINE_bool(
"directed", False, "Whether to generate directed graphs."
)
_OUTPUT_PATH = flags.DEFINE_string(
"output_path", None, "The output path to write the graphs.", required=True
)
_SPLIT = flags.DEFINE_string(
"split", None, "The dataset split to generate.", required=True
)
_MIN_SPARSITY = flags.DEFINE_float("min_sparsity", 0.0, "The minimum sparsity.")
_MAX_SPARSITY = flags.DEFINE_float("max_sparsity", 1.0, "The maximum sparsity.")
def write_graphs(graphs: list[nx.Graph], output_dir: str) -> None:
if not os.path.isdir(output_dir):
os.makedirs(output_dir)
for ind, graph in enumerate(graphs):
nx.write_graphml(
graph,
os.Open(
os.path.join(output_dir, str(ind) + ".graphml"),
"wb",
),
)
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError("Too many command-line arguments.")
if _SPLIT.value == "train":
random_seed = 9876
elif _SPLIT.value == "test":
random_seed = 1234
elif _SPLIT.value == "validation":
random_seed = 5432
else:
raise NotImplementedError()
generated_graphs = graph_generators.generate_graphs(
number_of_graphs=_NUMBER_OF_GRAPHS.value,
algorithm=_ALGORITHM.value,
directed=_DIRECTED.value,
random_seed=random_seed,
er_min_sparsity=_MIN_SPARSITY.value,
er_max_sparsity=_MAX_SPARSITY.value,
)
write_graphs(
graphs=generated_graphs,
output_dir=os.path.join(
_OUTPUT_PATH.value,
"directed" if _DIRECTED.value else "undirected",
_ALGORITHM.value,
_SPLIT.value,
),
)
if __name__ == "__main__":
app.run(main)
================================================
FILE: talk_like_a_graph/graph_generators_test.py
================================================
from absl.testing import parameterized
from . import graph_generators
from absl.testing import absltest
class GraphGenerationTest(absltest.TestCase, parameterized.TestCase):
@parameterized.named_parameters(
dict(
testcase_name='er_undirected_1',
algorithm='er',
directed=False,
k=1,
),
dict(
testcase_name='er_directed_1',
algorithm='er',
directed=True,
k=1,
),
dict(
testcase_name='ba_undirected_5',
algorithm='ba',
directed=False,
k=5,
),
dict(
testcase_name='ba_directed_5',
algorithm='ba',
directed=True,
k=5,
),
)
def test_number_of_graphs(self, algorithm, directed, k):
generated_graph = graph_generators.generate_graphs(k, algorithm, directed)
self.assertLen(generated_graph, k)
@parameterized.named_parameters(
dict(
testcase_name='er_undirected',
algorithm='er',
directed=False,
),
dict(
testcase_name='er_directed',
algorithm='er',
directed=True,
),
dict(
testcase_name='ba_undirected',
algorithm='ba',
directed=False,
),
dict(
testcase_name='ba_directed',
algorithm='ba',
directed=True,
),
dict(
testcase_name='sbm_undirected',
algorithm='sbm',
directed=False,
),
dict(
testcase_name='sbm_directed',
algorithm='sbm',
directed=True,
),
dict(
testcase_name='sfn_undirected',
algorithm='sfn',
directed=False,
),
dict(
testcase_name='sfn_directed',
algorithm='sfn',
directed=True,
),
dict(
testcase_name='complete_undirected',
algorithm='complete',
directed=False,
),
dict(
testcase_name='complete_directed',
algorithm='complete',
directed=True,
),
dict(
testcase_name='star_undirected',
algorithm='star',
directed=False,
),
dict(
testcase_name='star_directed',
algorithm='star',
directed=True,
),
dict(
testcase_name='path_undirected',
algorithm='path',
directed=False,
),
dict(
testcase_name='path_directed',
algorithm='path',
directed=True,
),
)
def test_directions(self, algorithm, directed):
generated_graph = graph_generators.generate_graphs(1, algorithm, directed)
self.assertEqual(generated_graph[0].is_directed(), directed)
if __name__ == '__main__':
googletest.main()
================================================
FILE: talk_like_a_graph/graph_metrics.py
================================================
"""Metrics for seqio tasks over graph data.
This module contains definitions of metric_fns to be used for scoring
graph tasks from nlgraph and graphqa.
"""
from typing import Mapping, Sequence
def yes_no_accuracy(
targets: Sequence[str], predictions: Sequence[str]
) -> Mapping[str, float]:
"""Assesses the accuracy of LLM outputs on Yes/No tasks.
Targets must contain either the word 'yes' or the word 'no' but not both.
Predictions are binarized by checking for 'yes' or 'no' in the first line.
Args:
targets: The expected output strings.
predictions: The LLM outputs.
Returns:
Returns a dict of the following metrics:
yes_no_accuracy: The % where the target and prediction match.
yes_no_ambiguous: The % where the prediction contained yes and no
yes_no_indeterminate: The % where the prediction contained neither yes nor
no
Raises:
ValueError: If a target string contains 'yes' and 'no'
"""
num_correct = 0
num_ambiguous = 0
num_indeterminate = 0
for target, prediction in zip(targets, predictions):
normalized_target = target.lower()
binarized_target = 'yes' in normalized_target
print(binarized_target)
if binarized_target and 'no' in normalized_target:
raise ValueError(f'Ambiguous target string, {target}')
if not binarized_target and 'no' not in normalized_target:
raise ValueError(f'Indeterminate target string, {target}')
normalized_prediction = prediction.splitlines()
if not normalized_prediction:
normalized_prediction = ''
else:
normalized_prediction = normalized_prediction[0]
normalized_prediction = normalized_prediction.lower()
binarized_prediction = 'yes' in normalized_prediction.lower()
print(normalized_prediction)
if binarized_prediction and 'no' in normalized_prediction:
num_ambiguous += 1
continue
if not binarized_prediction and 'no' not in normalized_prediction:
num_indeterminate += 1
continue
if binarized_prediction == binarized_target:
num_correct += 1
return {
'yes_no_accuracy': num_correct / len(targets),
'yes_no_ambiguous': num_ambiguous / len(targets),
'yes_no_indeterminate': num_indeterminate / len(targets),
}
================================================
FILE: talk_like_a_graph/graph_metrics_test.py
================================================
from . import graph_metrics
from absl.testing import absltest
class GraphTasksTest(absltest.TestCase):
def test_yes_no_correct(self):
result = graph_metrics.yes_no_accuracy(
targets=["Yes", "The answer is yes.", "No", "The answer is no.", "No"],
predictions=[
"yes",
"Yes\nNo",
"No\nYes",
"That's gonna be no from me.",
"yes",
],
)
self.assertEqual(result["yes_no_ambiguous"], 0)
self.assertEqual(result["yes_no_indeterminate"], 0)
self.assertAlmostEqual(result["yes_no_accuracy"], 0.8)
def test_yes_no_ambiguous(self):
result = graph_metrics.yes_no_accuracy(
targets=["Yes", "The answer is yes.", "No", "The answer is no.", "No"],
predictions=[
"yes",
"Yes No",
"No Yes",
"That's gonna be no from me.",
"yes",
],
)
self.assertEqual(result["yes_no_ambiguous"], 0.4)
self.assertEqual(result["yes_no_indeterminate"], 0)
self.assertAlmostEqual(result["yes_no_accuracy"], 0.4)
def test_yes_no_indeterminate(self):
result = graph_metrics.yes_no_accuracy(
targets=["Yes", "The answer is yes.", "No", "The answer is no.", "No"],
predictions=[
"yes",
"\n No",
"",
"That's gonna be no from me.",
"yes",
],
)
self.assertEqual(result["yes_no_ambiguous"], 0)
self.assertEqual(result["yes_no_indeterminate"], 0.4)
self.assertAlmostEqual(result["yes_no_accuracy"], 0.4)
def test_yes_no_accuracy_raises_on_ambiguous_target(self):
with self.assertRaises(ValueError):
graph_metrics.yes_no_accuracy(
targets=["Yes but maybe no"],
predictions=["yes"],
)
def test_yes_no_accuracy_raises_on_indeterminate_target(self):
with self.assertRaises(ValueError):
graph_metrics.yes_no_accuracy(
targets=["Hmm?"],
predictions=["yes"],
)
if __name__ == "__main__":
googletest.main()
================================================
FILE: talk_like_a_graph/graph_tasks.py
================================================
"""The graph tasks to be tried with LLMs."""
import random
import networkx as nx
import numpy as np
from . import graph_text_encoders
class GraphTask:
"""The parent class for all the graph tasks."""
def __init__(self):
self.name = 'default'
self.maximum_nnodes_cot_graph = 10
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
raise NotImplementedError()
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
):
raise NotImplementedError()
class CycleCheck(GraphTask):
"""The graph task to check if there is at least one cycle or not."""
def __init__(self):
super().__init__()
self.name = 'cycle_check'
self._task_description = 'Q: Is there a cycle in this graph?\nA: '
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
for ind, graph in enumerate(graphs):
question = (
graph_text_encoders.encode_graph(graph, encoding_method)
+ self._task_description
)
try:
nx.find_cycle(graph)
answer = 'Yes, there is a cycle.'
except nx.NetworkXNoCycle:
answer = 'No, there is no cycle.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': self._task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
"""Create a few shot example w or w/o cot for the graph graph."""
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = (
graph_text_encoders.encode_graph(graph, encoding_method)
+ self._task_description
)
try:
cycle = nx.find_cycle(graph)
cycle_text = ''
answer = 'Yes, there is a cycle. '
if cot:
for pair in cycle:
cycle_text += (
name_dict[pair[0]]
+ ' is connected to '
+ name_dict[pair[1]]
+ ', '
)
cycle_cot = 'The cycle is: %s.' % cycle_text[:-2]
answer += cycle_cot
except nx.NetworkXNoCycle:
answer = 'No, there is no cycle.'
return question + answer
def choose_few_shot_examples(
self,
few_shots_dict: dict[tuple[str, str], list[str]],
encoding_method: str,
k: int = 2,
) -> str:
"""Choose few shot examples for each algorithm."""
pos_cycle_algorithms = ['er', 'ba', 'sbm', 'sfn', 'complete']
neg_cycle_algorithms = ['star', 'path']
few_shots_str = ''
# choose k-1 shots for pos algorithms and one negative.
positive_algorithms = random.choices(pos_cycle_algorithms, k=k - 1)
for positive_algorithm in positive_algorithms:
example_list = few_shots_dict[(positive_algorithm, encoding_method)]
few_shots_str += 'Example: ' + random.choice(example_list) + '\n'
negative_algorithm = random.choice(neg_cycle_algorithms)
example_list = few_shots_dict[(negative_algorithm, encoding_method)]
few_shots_str += 'Example: ' + random.choice(example_list) + '\n'
return few_shots_str
class EdgeExistence(GraphTask):
"""The graph task to check if an edge exist in a graph or not."""
def __init__(self):
super().__init__()
self.name = 'edge_existence'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
task_description = 'Q: Is node %s connected to node %s?\nA: ' % (
name_dict[source],
name_dict[target],
)
question += task_description
if ((source, target) in graph.edges()) or (
(target, source) in graph.edges()
):
answer = 'Yes.'
else:
answer = 'No.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source, target],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += 'Q: Is node %s connected to node %s?\nA: ' % (
name_dict[source],
name_dict[target],
)
if ((source, target) in graph.edges()) or (
(target, source) in graph.edges()
):
answer = 'Yes.'
if cot:
answer += (
' Because, there is an edge from %s to %s in the graph description.'
% (name_dict[source], name_dict[target])
)
else:
answer = 'No.'
if cot:
answer += (
' Because, there is no edge from %s to %s in the graph description.'
% (name_dict[source], name_dict[target])
)
return question + answer
class NodeCount(GraphTask):
"""The graph task for finding number of nodes in a graph."""
def __init__(self):
super().__init__()
self.name = 'node_count'
self._task_description = 'Q: How many nodes are in this graph?\nA: '
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += self._task_description
answer = ' %d.' % len(graph.nodes())
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': self._task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [],
}
return examples_dict
def get_nodes_string(self, name_dict: dict[int, str], nnodes: int) -> str:
node_string = ''
for i in range(nnodes - 1):
node_string += name_dict[i] + ', '
node_string += 'and ' + name_dict[nnodes - 1]
return node_string
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += self._task_description
answer = '%d.' % len(graph.nodes())
if cot:
answer += ' The nodes are %s.' % self.get_nodes_string(
name_dict, len(graph.nodes())
)
return question + answer
class NodeDegree(GraphTask):
"""The graph task for finding degree of a node in a graph."""
def __init__(self):
super().__init__()
self.name = 'node_degree'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
task_description = (
'Q: What is the degree of node %s?\nA: ' % name_dict[source_node]
)
question += task_description
answer = '%d.' % graph.degree[source_node]
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source_node],
}
return examples_dict
def get_edge_string(
self, name_dict: dict[int, str], graph: nx.Graph, source_node: int
) -> str:
"""Gets a string identifying the edges a given node is connected to."""
edge_string = ''
target_edges = graph.edges(source_node)
target_nodes = []
for edge in target_edges:
target_nodes.append(edge[1])
if target_nodes:
for i in range(len(target_nodes) - 1):
edge_string += name_dict[target_nodes[i]] + ', '
edge_string += 'and ' + name_dict[target_nodes[-1]]
else:
edge_string = 'no nodes'
return edge_string
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
question += (
'Q: What is the degree of node %s?\nA: ' % name_dict[source_node]
)
answer = '%d.' % graph.degree[source_node]
if cot:
answer += ' This is because %s is connected to %s.' % (
name_dict[source_node],
self.get_edge_string(name_dict, graph, source_node),
)
return question + answer
class EdgeCount(GraphTask):
"""The graph task for finding number of edges in a graph."""
def __init__(self):
super().__init__()
self.name = 'edge_count'
self._task_description = 'Q: How many edges are in this graph?\nA: '
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += self._task_description
answer = ' %d.' % len(graph.edges())
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': self._task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [],
}
return examples_dict
def get_edges_string(
self, name_dict: dict[int, str], edges: list[tuple[int, int]]
) -> str:
edges_string = ''
for edge in edges:
edges_string += (
'(' + name_dict[edge[0]] + ', ' + name_dict[edge[1]] + '), '
)
return edges_string.strip()[:-1]
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += self._task_description
answer = '%d.' % len(graph.edges())
if cot:
answer += ' The edges are %s.' % self.get_edges_string(
name_dict, list(graph.edges())
)
return question + answer
class ConnectedNodes(GraphTask):
"""The graph task for finding connected nodes to a given node in a graph."""
def __init__(self):
super().__init__()
self.name = 'connected_nodes'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
task_description = (
'Q: List all the nodes connected to %s in alphabetical order.\nA: '
% name_dict[source_node]
)
question += task_description
outgoing_edges = list(graph.edges(source_node))
if outgoing_edges:
answer = self.get_connected_nodes(outgoing_edges, name_dict) + '.'
else:
answer = ' No nodes.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source_node],
}
return examples_dict
def get_connected_nodes(
self, edges: list[tuple[int, int]], name_dict: dict[int, str]
) -> str:
"""Gets a string including all the nodes that are connected to source."""
connected_nodes = []
for edge in edges:
connected_nodes.append(name_dict[edge[1]])
connected_nodes_string = ''
if connected_nodes:
try:
int(connected_nodes[0])
connected_nodes_string = ', '.join(map(str, connected_nodes))
except ValueError:
# Check if these are not integers, sort
connected_nodes_string = ', '.join(map(str, sorted(connected_nodes)))
return connected_nodes_string
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
question += (
'Q: List all the nodes connected to %s in alphabetical order.\nA: '
% name_dict[source_node]
)
outgoing_edges = list(graph.edges(source_node))
answer = ''
if outgoing_edges:
answer = self.get_connected_nodes(outgoing_edges, name_dict) + '.'
if cot:
answer += ' This is because there is an edge from %s to %s.' % (
name_dict[source_node],
answer,
)
else:
answer = 'No nodes.'
if cot:
answer += (
' This is because %s is not connected to any node.'
% name_dict[source_node]
)
return question + answer
class DisconnectedNodes(GraphTask):
"""The task for finding disconnected nodes for a given node in a graph."""
def __init__(self):
super().__init__()
self.name = 'disconnected_nodes'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
task_description = (
'Q: List all the nodes that are not connected to %s in alphabetical'
' order.\nA: '
% name_dict[source_node]
)
question += task_description
outgoing_edges = list(graph.edges(source_node))
answer = self.get_disconnected_nodes(
source_node, outgoing_edges, name_dict, list(graph.nodes())
)
if not answer:
answer = 'No nodes'
answer += '.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source_node],
}
return examples_dict
def get_disconnected_nodes(
self,
source: int,
edges: list[tuple[int, int]],
name_dict: dict[int, str],
all_nodes: list[int],
) -> str:
"""Gets a string with all the nodes that are not connected to source."""
for edge in edges:
if edge[1] in all_nodes:
all_nodes.remove(edge[1])
if source in all_nodes:
all_nodes.remove(source)
all_nodes_names = []
for node in all_nodes:
all_nodes_names.append(name_dict[node])
# sorted operation should be different for integers vs strings.
if all_nodes_names:
try:
int(all_nodes_names[0])
for ind, value in enumerate(all_nodes_names):
all_nodes_names[ind] = int(value)
all_nodes_names = sorted(all_nodes_names)
for ind, value in enumerate(all_nodes_names):
all_nodes_names[ind] = str(value)
except ValueError:
pass
return ', '.join(map(str, sorted(all_nodes_names)))
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
source_node = random.sample(list(graph.nodes()), k=1)[0]
question += (
'Q: List all the nodes that are not connected to %s in alphabetical'
' order.\nA: '
% name_dict[source_node]
)
outgoing_edges = list(graph.edges(source_node))
answer = ''
disconnected_nodes_string = self.get_disconnected_nodes(
source_node, outgoing_edges, name_dict, list(graph.nodes())
)
if outgoing_edges:
if not disconnected_nodes_string:
disconnected_nodes_string = 'No nodes'
answer = disconnected_nodes_string + '.'
if cot:
answer += ' This is because there is not an edge from %s to %s.' % (
name_dict[source_node],
answer,
)
else:
answer = ' No nodes.'
if cot:
answer += (
' This is because %s is connected to all nodes.'
% name_dict[source_node]
)
return question + answer
class Reachability(GraphTask):
"""The graph task to check if there is a path from a source to target."""
def __init__(self):
super().__init__()
self.name = 'reachability'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
task_description = 'Q: Is there a path from node %s to node %s?\nA: ' % (
name_dict[source],
name_dict[target],
)
question += task_description
if nx.has_path(graph, source, target):
answer = 'Yes.'
else:
answer = 'No.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source, target],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += 'Q: Is there a path from node %s to node %s?\nA: ' % (
name_dict[source],
name_dict[target],
)
if nx.has_path(graph, source, target):
answer = 'Yes.'
if cot:
path = nx.shortest_path(graph, source, target)
explanation = ' Because'
for i in range(len(path) - 1):
# The only edge or the non-last edges in the path.
if len(path) == 2 or i < len(path) - 2:
sep = ','
# The last edge in a path with more than one edge.
else:
sep = ', and'
explanation += '%s there is an edge from node %d to node %d' % (
sep,
path[i],
path[i + 1],
)
explanation += ' .'
answer += explanation
else:
answer = 'No.'
if cot:
answer += (
' Because, there is no path connecting node %s to node %s based on'
' the graph description.' % (name_dict[source], name_dict[target])
)
return question + answer
class ShortestPath(GraphTask):
"""The graph task to check if there is a path from a source to target."""
def __init__(self):
super().__init__()
self.name = 'shortest_path'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
task_description = (
'Q: What is the length of the shortest path from node %s to node'
' %s?\nA: '
% (
name_dict[source],
name_dict[target],
)
)
question += task_description
try:
path = nx.shortest_path(graph, source, target)
answer = str(len(path) - 1) + '.'
except nx.NetworkXNoPath:
answer = 'There is no path from node %s to node %s.' % (
name_dict[source],
name_dict[target],
)
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source, target],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += (
'Q: What is the length of the shortest path from node %s to node'
' %s?\nA: '
% (
name_dict[source],
name_dict[target],
)
)
if nx.has_path(graph, source, target):
path = nx.shortest_path(graph, source, target)
answer = str(len(path) - 1) + '.'
if cot:
explanation = ' Because'
for i in range(len(path) - 1):
# The only edge or the non-last edges in the path.
if len(path) == 2 or i < len(path) - 2:
sep = ','
# The last edge in a path with more than one edge.
else:
sep = ', and'
explanation += '%s there is an edge from node %d to node %d' % (
sep,
path[i],
path[i + 1],
)
explanation += ' .'
answer += explanation
else:
answer = 'There is no path from node %s to node %s.' % (
name_dict[source],
name_dict[target],
)
if cot:
answer += (
' Because, there is no path connecting node %s to node %s based on'
' the graph description.' % (name_dict[source], name_dict[target])
)
return question + answer
class TriangleCounting(GraphTask):
"""The graph task to count the number of triangles in a graph."""
def __init__(self):
super().__init__()
self.name = 'triangle_counting'
self._task_description = 'Q: How many triangles are in this graph?\nA: '
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
for ind, graph in enumerate(graphs):
question = (
graph_text_encoders.encode_graph(graph, encoding_method)
+ self._task_description
)
ntriangles = int(np.sum(list(nx.triangles(graph).values())) / 3)
answer = '%i.' % ntriangles
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': self._task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
"""Create a few shot example w or w/o cot for the graph graph."""
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = (
graph_text_encoders.encode_graph(graph, encoding_method)
+ self._task_description
)
triangles_dict = nx.triangles(graph)
ntriangles = int(np.sum(list(triangles_dict.values())) / 3)
if ntriangles > 0:
answer = '%i.' % ntriangles
if cot:
ntriangles_cot = ''
for key, value in triangles_dict.items():
if value > 0:
if value == 1:
ntriangles_cot += (
'There is %i triangle including node %s as a vertex.\n'
% (value, name_dict[key])
)
else:
ntriangles_cot += (
'There are %i triangles including node %s as a vertex.\n'
% (value, name_dict[key])
)
ntriangles_cot += (
'Summing the number of triangles for all nodes and dividing them by'
' three gives us %i triangles in total.' % ntriangles
)
answer += ntriangles_cot
else:
answer = '0.'
if cot:
ntriangles_cot = 'No three nodes form a triangle of edges.'
answer += ntriangles_cot
return question + answer
class MaximumFlow(GraphTask):
"""The graph task to compute the maximum flow from a source to a target."""
def __init__(self):
super().__init__()
self.name = 'maximum_flow'
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
graph = add_edge_weight(graph)
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
task_description = (
'Q: What is the maximum capacity of the flow from node %s to node'
' %s?\nA: ' % (name_dict[source], name_dict[target])
)
question += task_description
maximum_flow_value = nx.maximum_flow(
graph, source, target, capacity='weight'
)[0]
answer = str(maximum_flow_value) + '.'
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(len(graph.nodes())),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
'node_ids': [source, target],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
graph = add_edge_weight(graph)
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
source, target = random.sample(list(graph.nodes()), k=2)
question = graph_text_encoders.encode_graph(graph, encoding_method)
question += (
'Q: What is the maximum capacity of the flow from node %s to'
' node %s?\nA: ' % (name_dict[source], name_dict[target])
)
flow_value, flow_dict = nx.maximum_flow(
graph, source, target, capacity='weight'
)
answer = str(flow_value) + '.'
if flow_value > 0:
if cot:
explanation = ' This is because of the following edges: '
for edge, capacity in flow_dict.items():
for key, value in capacity.items():
if value > 0:
explanation += (
'the edge from node %i to node %i with capacity %i, '
% (
edge,
key,
value,
)
)
explanation = explanation.strip()[:-1] + '.'
answer += explanation
else:
if cot:
answer += (
' Because, there is no path connecting node %s to node %s based on'
' the graph description.' % (name_dict[source], name_dict[target])
)
return question + answer
def has_edge_weights(graph):
for _, _, data in graph.edges(data=True):
if 'weight' not in data:
return False
return True
def add_edge_weight(graph):
if has_edge_weights(graph):
return graph
else:
for edge in graph.edges():
graph[edge[0]][edge[1]]['weight'] = random.randint(1, 10)
return graph
class NodeClassification(GraphTask):
"""The graph task to classify a given node in the graph."""
def __init__(self):
super().__init__()
self.name = 'node_classification'
self.classes = [
'soccer',
'baseball',
'tennis',
'golf',
'football',
'surfing',
]
def prepare_examples_dict(
self,
graphs: list[nx.Graph],
generator_algorithms: list[str],
encoding_method: str,
) -> dict[int, dict[str, str | list[int]]]:
classes = random.sample(list(self.classes), k=2)
examples_dict = {}
name_dict = graph_text_encoders.get_tlag_node_encoder(None, encoding_method)
for ind, graph in enumerate(graphs):
question = graph_text_encoders.encode_graph(graph, encoding_method)
nnodes = len(graph.nodes())
# Sampling nnodes // 2 + 1 nodes.
sampled_nodes = random.sample(
list(graph.nodes(data=True)), k=nnodes // 2 + 1
)
# Adding the class of half of the nodes.
for node_data in sampled_nodes[:-1]:
node_class = classes[node_data[1]['block']]
question += (
'Node ' + name_dict[node_data[0]] + ' likes ' + node_class + '.\n'
)
# Reserving the last sampled node for the question.
task_description = 'Q: Does node %s like %s or %s?\nA: ' % (
name_dict[sampled_nodes[-1][0]],
classes[0],
classes[1],
)
question += task_description
answer = classes[sampled_nodes[-1][1]['block']]
examples_dict[ind] = {
'question': question,
'answer': answer,
'nnodes': str(nnodes),
'nedges': str(len(graph.edges())),
'task_description': task_description,
'graph': graph,
'algorithm': generator_algorithms[ind],
# id of the last samples node
'node_ids': [sampled_nodes[-1][0]],
}
return examples_dict
def create_few_shot_example(
self, graph: nx.Graph, encoding_method: str, cot: bool
) -> str:
classes = random.sample(list(self.classes), k=2)
name_dict = graph_text_encoders.get_tlag_node_encoder(
graph, encoding_method
)
question = graph_text_encoders.encode_graph(graph, encoding_method)
nnodes = len(graph.nodes())
sampled_nodes = random.sample(
list(graph.nodes(data=True)), k=nnodes // 2 + 1
)
for node_data in sampled_nodes[:-1]:
node_class = classes[node_data[1]['block']]
question += (
'Node ' + name_dict[node_data[0]] + ' likes ' + node_class + '.\n'
)
task_description = 'Q: Does node %s like %s or %s?\nA: ' % (
name_dict[sampled_nodes[-1][0]],
classes[0],
classes[1],
)
question += task_description
answer = classes[sampled_nodes[-1][1]['block']]
if cot:
explanation = (
' This is because most of the nodes that are connected to node %s'
' likes %s.'
% (sampled_nodes[-1][0], classes[sampled_nodes[-1][1]['block']])
)
answer += explanation
return question + answer
================================================
FILE: talk_like_a_graph/graph_tasks_generator.py
================================================
r"""The graph tasks to be tried with LLMs..
This code loads graphs and creates graph tasks and output them as tf examples in
a recordio file in the task directory provided.
# Placeholder for Google-internal comments.
"""
from collections.abc import Sequence
import os
import random
from absl import app
from absl import flags
import networkx as nx
import numpy as np
from . import graph_tasks
from . import graph_tasks_utils as utils
_TASK_DIR = flags.DEFINE_string(
'task_dir', None, 'The directory to write tasks.', required=True
)
_GRAPHS_DIR = flags.DEFINE_string(
'graphs_dir', None, 'The directory containing the graphs.', required=True
)
_RANDOM_SEED = flags.DEFINE_integer(
'random_seed',
None,
'The random seed to use for task generation.',
required=True,
)
def zero_shot(
task: graph_tasks.GraphTask,
graphs: list[nx.Graph],
algorithms: list[str],
text_encoders: list[str],
cot: bool,
random_seed: int,
split: str,
) -> None:
"""Creating zero-shot or zero-cot examples for the given task.
Args:
task: the corresponding graph task.
graphs: the list of graphs to use for the task.
algorithms: the algorithm used to generate the graphs.
text_encoders: the encoders to use in the tasks.
cot: whether to apply cot or not.
random_seed: the random seed to use in the process.
split: whether we are creating a train or test split.
"""
random.seed(random_seed)
zero_shot_examples = utils.create_zero_shot_task(
task, graphs, algorithms, text_encoders, cot=cot
)
file_name = task.name + ('_zero_cot_' if cot else '_zero_shot_')
file_name += split + '.recordio'
utils.write_examples(
zero_shot_examples,
os.path.join(_TASK_DIR.value, file_name),
)
def few_shot(
task: graph_tasks.GraphTask,
graphs: list[nx.Graph],
few_shot_graphs: list[nx.Graph],
algorithms: list[str],
text_encoders: list[str],
cot: bool,
bag: bool,
random_seed: int,
) -> None:
"""Creating few-shot, cot, or cot-bag examples for the given task.
Args:
task: the corresponding graph task.
graphs: the list of graphs to use for the task.
few_shot_graphs: the list of graphs to generate few shot examples for.
algorithms: the algorithm used to generate the graphs.
text_encoders: the encoders to use in the tasks.
cot: whether to apply cot or not.
bag: whether to apply build-a-graph method or not.
random_seed: the random seed to use in the process.
"""
random.seed(random_seed)
few_shot_examples = utils.create_few_shot_task(
task,
graphs,
algorithms,
few_shot_graphs,
text_encoders,
cot=cot,
bag=bag,
random_seed=random_seed,
)
file_name = task.name
if cot and bag:
file_name += '_few_shot_cot_bag_test.recordio'
elif cot:
file_name += '_few_shot_cot_test.recordio'
else:
file_name += '_few_shot_test.recordio'
utils.write_examples(
few_shot_examples,
os.path.join(_TASK_DIR.value, file_name),
)
def generate_random_sbm_graph(random_state: np.random.RandomState):
# Sampling a small number as the probability of the two nodes in different
# communities being connected.
small_number = random.uniform(0, 0.05)
# Sampling a large number as probability of the nodes in one community
# being connected.
large_number = random.uniform(0.6, 0.8)
number_of_nodes = random.choice(np.arange(5, 20))
sizes = [number_of_nodes // 2, number_of_nodes // 2]
probs = [[large_number, small_number], [small_number, large_number]]
return nx.stochastic_block_model(sizes, probs, seed=random_state)
def main(argv: Sequence[str]) -> None:
if len(argv) > 1:
raise app.UsageError('Too many command-line arguments.')
algorithms = ['er']
directions = ['undirected']
text_encoders = ['adjacency']
# Loading the graphs.
graphs = []
generator_algorithms = []
for algorithm in algorithms:
for direction in directions:
loaded_graphs = utils.load_graphs(
_GRAPHS_DIR.value,
algorithm,
'train',
direction,
)
graphs += loaded_graphs
generator_algorithms += [algorithm] * len(loaded_graphs)
# Defining a task on the graphs
task = graph_tasks.ShortestPath()
if isinstance(task, graph_tasks.NodeClassification):
# The node classification task requires SBM graphs. As it's not possible to
# write graphs with data (e.g., blocks data as in SBM graphs), we regenerate
# graphs.
random_state = np.random.RandomState(_RANDOM_SEED.value)
print('Generating sbm graphs')
graphs = [
generate_random_sbm_graph(random_state) for _ in range(len(graphs))
]
zero_shot(
task,
graphs,
generator_algorithms,
text_encoders,
cot=False,
random_seed=_RANDOM_SEED.value,
split='test',
)
zero_shot(
task,
graphs,
generator_algorithms,
text_encoders,
cot=True,
random_seed=_RANDOM_SEED.value,
split='test',
)
# Loading few-shot graphs.
few_shot_graphs = []
for algorithm in algorithms:
for direction in directions:
few_shot_graphs += utils.load_graphs(
_GRAPHS_DIR.value,
algorithm,
'train',
direction,
)
if isinstance(task, graph_tasks.NodeClassification):
# The node classification task requires SBM graphs. As it's not possible to
# write graphs with data (e.g., blocks data as in SBM graphs), we regenerate
# graphs.
random_state = np.random.RandomState(_RANDOM_SEED.value + 1)
print('Generating few shot sbm graphs')
few_shot_graphs = [
generate_random_sbm_graph(random_state)
for _ in range(len(few_shot_graphs))
]
few_shot(
task,
graphs,
few_shot_graphs,
generator_algorithms,
text_encoders,
cot=False,
bag=False,
random_seed=_RANDOM_SEED.value,
)
few_shot(
task,
graphs,
few_shot_graphs,
generator_algorithms,
text_encoders,
cot=True,
bag=False,
random_seed=_RANDOM_SEED.value,
)
few_shot(
task,
graphs,
few_shot_graphs,
generator_algorithms,
text_encoders,
cot=True,
bag=True,
random_seed=_RANDOM_SEED.value,
)
if __name__ == '__main__':
app.run(main)
================================================
FILE: talk_like_a_graph/graph_tasks_utils.py
================================================
"""The graph tasks to be tried with LLMs."""
import os
import random
import networkx as nx
import numpy as np
import seqio
import tensorflow as tf
import tensorflow_gnn as tfgnn
# Google-internal import(s).
# Internal import.
from . import graph_tasks
from tensorflow.core.example import example_pb2
from tensorflow.core.example import feature_pb2
def laplacian_pos_embedding(graph: nx.Graph, units: int = 4) -> nx.Graph:
"""Adds the laplacian positional encoding."""
m = nx.normalized_laplacian_matrix(
graph, nodelist=sorted(graph.nodes), weight=None
).astype(np.float32)
u, _, _ = np.linalg.svd(m.todense(), compute_uv=True)
if units > u.shape[1]:
u = np.pad(u, ((0, 0), (0, units - u.shape[1])))
nx.set_node_attributes(
graph, dict(zip(sorted(graph.nodes), u[:, :units])), name='lpe'
)
return graph
def to_tfgnn(graph: nx.Graph, node_ids: list[int]) -> tfgnn.GraphTensor:
"""Convert a given nx graph to a tfgnn graph."""
if graph.edges(data=True):
s, t, w = zip(*[
(s, t, (d['weight'] if d and 'weight' in d else None))
for s, t, d in graph.edges(data=True)
])
else:
s, t, w = (), (), ()
# tfgnn assumes graphs are directed. Adding the rev edges for an undirected
# graph.
if not graph.is_directed():
s, t, w = s + t, t + s, w + w
graph = laplacian_pos_embedding(graph, units=4)
features = set(k for n in graph.nodes for k in graph.nodes[n].keys()) # pylint: disable=g-complex-comprehension
node_features = {
f: tf.convert_to_tensor([graph.nodes[n][f] for n in graph.nodes])
for f in features
}
# If all edges have a non-trivial weight, then we record the weights.
if all(w):
edge_features = {'weights': tf.convert_to_tensor(w, dtype=tf.int32)}
gt = tfgnn.homogeneous(
tf.convert_to_tensor(s, dtype=tf.int32),
tf.convert_to_tensor(t, dtype=tf.int32),
node_features=node_features,
edge_features=edge_features,
)
else:
gt = tfgnn.homogeneous(
tf.convert_to_tensor(s, dtype=tf.int32),
tf.convert_to_tensor(t, dtype=tf.int32),
node_features=node_features,
)
if not node_ids:
# No node is mentioned in the task description.
return gt
node_sets = {
**gt.node_sets,
'_readout': tfgnn.NodeSet.from_fields(sizes=[1]),
}
if len(node_ids) == 1:
# Tasks requiring only one node id e.g., computing node degree.
edge_sets = {
**gt.edge_sets,
'_readout/node': tfgnn.EdgeSet.from_fields(
sizes=[1],
adjacency=tfgnn.Adjacency.from_indices(
source=('nodes', node_ids),
target=('_readout', [0]),
),
),
}
elif len(node_ids) == 2:
# Tasks requiring two nodes e.g., shortest path from one node to the other.
edge_sets = {
**gt.edge_sets,
'_readout/source': tfgnn.EdgeSet.from_fields(
sizes=[1],
adjacency=tfgnn.Adjacency.from_indices(
source=('nodes', node_ids[:1]),
target=('_readout', [0]),
),
),
'_readout/target': tfgnn.EdgeSet.from_fields(
sizes=[1],
adjacency=tfgnn.Adjacency.from_indices(
source=('nodes', node_ids[1:]),
target=('_readout', [0]),
),
),
}
else:
# Raising an error if more than two nodes are mentiones.
raise ValueError(f'Invalid number of integers: {len(node_ids)}')
return tfgnn.GraphTensor.from_pieces(
context=gt.context, node_sets=node_sets, edge_sets=edge_sets
)
def create_example_feature(
key: int,
question: str,
answer: str,
algorithm: str,
encoding_method: str,
nnodes: str,
nedges: str,
task_description: str,
graph: nx.Graph,
node_ids: list[int],
) -> example_pb2.Example:
"""Create a tensorflow example from a datapoint."""
key_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[str(key).encode()])
)
question_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[question.encode()])
)
answer_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[answer.encode()])
)
algorithm_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[algorithm.encode()])
)
encoding_method_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[encoding_method.encode()])
)
nnodes_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[nnodes.encode()])
)
nedges_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[nedges.encode()])
)
task_description_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[task_description.encode()])
)
gt = to_tfgnn(graph, node_ids)
graph_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(
value=[tfgnn.write_example(gt).SerializeToString()]
)
)
directed_feature = feature_pb2.Feature(
bytes_list=tf.train.BytesList(value=[str(graph.is_directed()).encode()])
)
example_feats = tf.train.Features(
feature={
'id': key_feature,
'question': question_feature,
'answer': answer_feature,
'algorithm': algorithm_feature,
'text_encoding': encoding_method_feature,
'nnodes': nnodes_feature,
'nedges': nedges_feature,
'task_description': task_description_feature,
'graph': graph_feature,
'directed': directed_feature,
}
)
return example_pb2.Example(features=example_feats)
def load_graphs(
base_path: str,
algorithm: str,
split: str,
direction: str,
max_nnodes: int = 20,
) -> list[nx.Graph]:
"""Load a list of graphs from a given algorithm and split."""
graphs_path = os.path.join(
base_path,
direction,
algorithm,
split,
)
loaded_graphs = []
all_files = gfile.ListDir(graphs_path)
for file in all_files:
if file.endswith('.graphml'):
path = os.path.join(graphs_path, file)
graph = nx.read_graphml(os.Open(path, 'rb'), node_type=int)
if graph.number_of_nodes() <= max_nnodes:
loaded_graphs.append(graph)
return loaded_graphs
def prepare_examples(
examples_dict: dict[int, dict[str, str | list[int]]],
encoding_method: str,
) -> list[example_pb2.Example]:
"""Create a list of tf.train.Example from a dict of examples."""
examples = []
for key, value in examples_dict.items():
(
question,
answer,
nnodes,
nedges,
task_description,
graph,
algorithm,
node_ids,
) = (
value['question'],
value['answer'],
value['nnodes'],
value['nedges'],
value['task_description'],
value['graph'],
value['algorithm'],
value['node_ids'],
)
examples.append(
create_example_feature(
key,
question,
answer,
algorithm,
encoding_method,
nnodes,
nedges,
task_description,
graph,
node_ids,
)
)
return examples
def create_zero_shot_task(
task: graph_tasks.GraphTask,
graphs: list[nx.Graph],
generator_algorithms: list[str],
text_encoders: list[str],
cot: bool = False,
) -> list[example_pb2.Example]:
"""Create a recordio file with zero-shot examples for the task."""
examples = []
for encoding_method in text_encoders:
examples_dict = task.prepare_examples_dict(
graphs, generator_algorithms, encoding_method
)
if cot:
for key in examples_dict.keys():
examples_dict[key]['question'] += "Let's think step by step. "
examples += prepare_examples(examples_dict, encoding_method)
return examples
def write_examples(examples: list[example_pb2.Example], output_path: str):
with recordio.RecordWriter(output_path) as output_file:
for example in examples:
output_file.WriteRecord(example.SerializeToString())
def prepare_few_shots(
task: graph_tasks.GraphTask,
graphs: list[nx.Graph],
text_encoders: list[str],
cot: bool,
) -> dict[str, list[str]]:
"""Create a dict of few-shot examples with their cot for the task."""
few_shots_examples_dict = {}
for encoding_method in text_encoders:
if encoding_method not in few_shots_examples_dict:
few_shots_examples_dict[(encoding_method)] = []
for graph in graphs:
few_shots_examples_dict[(encoding_method)].append(
task.create_few_shot_example(graph, encoding_method, cot)
)
return few_shots_examples_dict
def choose_few_shot_examples(
few_shots_dict: dict[str, list[str]],
encoding_method: str,
k: int = 2,
) -> str:
"""Choose few shot examples for each algorithm."""
few_shots_str = ''
for _ in range(k):
example_list = few_shots_dict[encoding_method]
few_shots_str += 'Example: ' + random.choice(example_list) + '\n'
return few_shots_str
def create_few_shot_task(
task: graph_tasks.GraphTask,
graphs: list[nx.Graph],
generator_algorithms: list[str],
few_shots_graphs: list[nx.Graph],
text_encoders: list[str],
cot: bool,
bag: bool,
random_seed: int,
) -> list[example_pb2.Example]:
"""Create a recordio file with few-shot examples for the task."""
# LINT.IfChange
vocab_path = None
# LINT.ThenChange(//research/graph/llm/graphqa/copy.bara.sky)
# Loading the palm tokenizer to calculate number of tokens in the sequence.
sp_vocab = seqio.SentencePieceVocabulary(vocab_path)
number_of_tokens = {}
examples = []
print('prepare few shot task', 'cot', cot, 'bag', bag)
few_shots_examples_dict = prepare_few_shots(
task,
few_shots_graphs,
text_encoders,
cot,
)
for encoding_method in text_encoders:
random.seed(random_seed)
examples_dict = task.prepare_examples_dict(
graphs, generator_algorithms, encoding_method
)
for key in examples_dict.keys():
few_shots_examples = choose_few_shot_examples(
few_shots_examples_dict,
encoding_method,
)
examples_dict[key]['question'] = (
few_shots_examples + 'Example: ' + examples_dict[key]['question']
)
if bag:
examples_dict[key]['question'] = examples_dict[key]['question'].replace(
'\nQ: ',
"\nLet's construct the graph with the nodes and edges first.\nQ: ",
) # pytype: disable=attribute-error
if encoding_method not in number_of_tokens:
number_of_tokens[encoding_method] = []
number_of_tokens[encoding_method].append(
len(sp_vocab.encode(examples_dict[key]['question']))
)
examples += prepare_examples(examples_dict, encoding_method)
# Printing maximum number of tokens in the sequence.
for key, value in number_of_tokens.items():
print(key, np.max(value))
return examples
================================================
FILE: talk_like_a_graph/graph_text_encoders.py
================================================
"""Library for encoding graphs in text."""
import networkx as nx
from . import name_dictionaries
def create_node_string(name_dict, nnodes: int) -> str:
node_string = ""
sorted_keys = list(sorted(name_dict.keys()))
for i in sorted_keys[: nnodes - 1]:
node_string += name_dict[i] + ", "
node_string += "and " + name_dict[sorted_keys[nnodes - 1]]
return node_string
def nx_encoder(graph: nx.Graph, _: dict[int, str], edge_type="id") -> str:
"""Encoding a graph as entries of an adjacency matrix."""
if graph.is_directed():
output = (
"In a directed graph, (s,p,o) means that there is an edge from node s"
" to node o of type p. "
)
else:
output = (
"In an undirected graph, (s,p,o) means that node s and node o are"
" connected with an undirected edge of type p. "
)
name_dict = {x: str(x) for x in graph.nodes()}
nodes_string = create_node_string(name_dict, nnodes=len(graph.nodes()))
output += "G describes a graph among nodes %s.\n" % nodes_string
if graph.edges():
output += "The edges in G are: "
for i, j in graph.edges():
edge_type = graph.get_edge_data(i, j)[edge_type]
if edge_type is None:
edge_type = "linked"
output += "(%s, %s, %s) " % (name_dict[i], edge_type, name_dict[j])
return output.strip() + ".\n"
def adjacency_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
"""Encoding a graph as entries of an adjacency matrix."""
if graph.is_directed():
output = (
"In a directed graph, (i,j) means that there is an edge from node i to"
" node j. "
)
else:
output = (
"In an undirected graph, (i,j) means that node i and node j are"
" connected with an undirected edge. "
)
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output += "G describes a graph among nodes %s.\n" % nodes_string
if graph.edges():
output += "The edges in G are: "
for i, j in graph.edges():
output += "(%s, %s) " % (name_dict[i], name_dict[j])
return output.strip() + ".\n"
def friendship_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
"""Encoding a graph as a friendship graph."""
if graph.is_directed():
raise ValueError("Friendship encoder is not defined for directed graphs.")
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output = (
"G describes a friendship graph among nodes %s.\n" % nodes_string.strip()
)
if graph.edges():
output += "We have the following edges in G:\n"
for i, j in graph.edges():
output += "%s and %s are friends.\n" % (name_dict[i], name_dict[j])
return output
def coauthorship_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
"""Encoding a graph as a coauthorship graph."""
if graph.is_directed():
raise ValueError("Coauthorship encoder is not defined for directed graphs.")
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output = (
"G describes a coauthorship graph among nodes %s.\n"
% nodes_string.strip()
)
if graph.edges():
output += "In this coauthorship graph:\n"
for i, j in graph.edges():
output += "%s and %s wrote a paper together.\n" % (
name_dict[i],
name_dict[j],
)
return output.strip() + ".\n"
def incident_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
"""Encoding a graph with its incident lists."""
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output = "G describes a graph among nodes %s.\n" % nodes_string
if graph.edges():
output += "In this graph:\n"
for source_node in graph.nodes():
target_nodes = graph.neighbors(source_node)
target_nodes_str = ""
nedges = 0
for target_node in target_nodes:
target_nodes_str += name_dict[target_node] + ", "
nedges += 1
if nedges > 1:
output += "Node %s is connected to nodes %s.\n" % (
source_node,
target_nodes_str[:-2],
)
elif nedges == 1:
output += "Node %d is connected to node %s.\n" % (
source_node,
target_nodes_str[:-2],
)
return output
def social_network_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
"""Encoding a graph as a social network graph."""
if graph.is_directed():
raise ValueError(
"Social network encoder is not defined for directed graphs."
)
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output = (
"G describes a social network graph among nodes %s.\n"
% nodes_string.strip()
)
if graph.edges():
output += "We have the following edges in G:\n"
for i, j in graph.edges():
output += "%s and %s are connected.\n" % (name_dict[i], name_dict[j])
return output
def expert_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
nodes_string = create_node_string(name_dict, len(graph.nodes()))
output = (
"You are a graph analyst and you have been given a graph G among nodes"
" %s.\n"
% nodes_string.strip()
)
output += "G has the following undirected edges:\n" if graph.edges() else ""
for i, j in graph.edges():
output += "%s -> %s\n" % (name_dict[i], name_dict[j])
return output
def nodes_to_text(graph, encoding_type):
"""Get dictionary converting node ids to text."""
if encoding_type == "integer":
return name_dictionaries.create_name_dict(graph, "integer", nnodes=1000)
elif encoding_type == "popular":
return name_dictionaries.create_name_dict(graph, "popular")
elif encoding_type == "alphabet":
return name_dictionaries.create_name_dict(graph, "alphabet")
elif encoding_type == "got":
return name_dictionaries.create_name_dict(graph, "got")
elif encoding_type == "south_park":
return name_dictionaries.create_name_dict(graph, "south_park")
elif encoding_type == "politician":
return name_dictionaries.create_name_dict(graph, "politician")
elif encoding_type == "random":
return name_dictionaries.create_name_dict(
graph, "random_integer", nnodes=1000
)
elif encoding_type == "nx_node_name":
return name_dictionaries.create_name_dict(graph, "nx_node_name")
else:
raise ValueError("Unknown encoding type: %s" % encoding_type)
def get_tlag_node_encoder(graph, encoder_name):
"""Find the node encoder used in the 'Talk Like a Graph' paper."""
if encoder_name == "adjacency":
return nodes_to_text(graph, "integer")
elif encoder_name == "incident":
return nodes_to_text(graph, "integer")
elif encoder_name == "friendship":
return nodes_to_text(graph, "popular")
elif encoder_name == "south_park":
return nodes_to_text(graph, "south_park")
elif encoder_name == "got":
return nodes_to_text(graph, "got")
elif encoder_name == "politician":
return nodes_to_text(graph, "politician")
elif encoder_name == "social_network":
return nodes_to_text(graph, "popular")
elif encoder_name == "expert":
return nodes_to_text(graph, "expert")
elif encoder_name == "coauthorship":
return nodes_to_text(graph, "popular")
elif encoder_name == "random":
return nodes_to_text(graph, "random")
elif encoder_name == "nx_node_name":
return nodes_to_text(graph, "nx_node_name")
else:
raise ValueError("Unknown graph encoder strategy: %s" % encoder_name)
# A dictionary from edge encoder name to the corresponding function.
EDGE_ENCODER_FN = {
"adjacency": adjacency_encoder,
"incident": incident_encoder,
"friendship": friendship_encoder,
"south_park": friendship_encoder,
"got": friendship_encoder,
"politician": social_network_encoder,
"social_network": social_network_encoder,
"expert": expert_encoder,
"coauthorship": coauthorship_encoder,
"random": adjacency_encoder,
"nx_edge_encoder": nx_encoder,
}
def with_ids(graph: nx.Graph, node_encoder: str) -> nx.Graph:
nx.set_node_attributes(graph, nodes_to_text(graph, node_encoder), name="id")
return graph
def encode_graph(
graph: nx.Graph, graph_encoder=None, node_encoder=None, edge_encoder=None
) -> str:
r"""Encodes a graph as text.
This relies on choosing:
a node_encoder and an edge_encoder:
or
a graph_encoder (a predefined pair of node and edge encoding strategies).
Note that graph_encoders may assume that the graph has some properties
(e.g. integer keys).
Example usage:
.. code-block:: python
```
# Use a predefined graph encoder from the paper.
>>> G = nx.karate_club_graph()
>>> encode_graph(G, graph_encoder="adjacency")
'In an undirected graph, (i,j) means that node i and node j are
connected
with an undirected edge. G describes a graph among nodes 0, 1, 2, 3, 4, 5,
6,
7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26,
27, 28, 29, 30, 31, 32, and 33.\nThe edges in G are: (0, 1) (0, 2) (0, 3)
...'
# Use the node's name in the graph as the node identifier.
>>> G = nx.les_miserables_graph()
>>> encode_graph(G, node_encoder="nx_node_name", edge_encoder="friendship")
'G describes a friendship graph among nodes Anzelma, Babet, Bahorel,
Bamatabois, BaronessT, Blacheville, Bossuet, Boulatruelle, Brevet, ...
We have the following edges in G:
Napoleon and Myriel are friends. Myriel and MlleBaptistine are friends...'
# Use the `id` feature from the edges to describe the edge type.
>>> G = nx.karate_club_graph()
>>> encode_graph(G, node_encoder="nx_node_name", edge_encoder="nx_edge_id")
'In an undirected graph, (s,p,o) means that node s and node o are connected
with an undirected edge of type p. G describes a graph among nodes 0, 1, 2, 3,
4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24,
25, 26, 27, 28, 29, 30, 31, 32, and 33.
The edges in G are: (0, linked, 1) (0, linked, 2) (0, linked, 3) ...'
```
Args:
graph: the graph to be encoded.
graph_encoder: the name of the graph encoder to use.
node_encoder: the name of the node encoder to use.
edge_encoder: the name of the edge encoder to use.
Returns:
The encoded graph as a string.
"""
# Check that only one of graph_encoder or (node_encoder, edge_encoder) is set.
if graph_encoder and (node_encoder or edge_encoder):
raise ValueError(
"Only one of graph_encoder or (node_encoder, edge_encoder) can be set."
)
if graph_encoder:
if isinstance(graph_encoder, str):
node_encoder_dict = get_tlag_node_encoder(graph, graph_encoder)
return EDGE_ENCODER_FN[graph_encoder](graph, node_encoder_dict)
else:
return graph_encoder(graph)
else:
node_encoder_dict = nodes_to_text(graph, node_encoder)
return EDGE_ENCODER_FN[edge_encoder](graph, node_encoder_dict)
================================================
FILE: talk_like_a_graph/graph_text_encoders_test.py
================================================
"""Testing for graph_text_encoders.py."""
from absl.testing import parameterized
import networkx as nx
from . import graph_text_encoders
from absl.testing import absltest
_G = nx.Graph()
_G.add_node(0)
_G.add_node(1)
_G.add_node(2)
_G.add_node(3)
_G.add_edge(0, 1)
_G.add_edge(1, 2)
_G.add_edge(2, 3)
_G.add_edge(3, 0)
class GraphTextEncodersTest(absltest.TestCase, parameterized.TestCase):
@parameterized.named_parameters(
dict(
testcase_name='adjacency_integer',
encoding_method='adjacency',
expected_result=(
'In an undirected graph, (i,j) means that node i and node j are'
' connected with an undirected edge. G describes a graph among'
' nodes 0, 1, 2, and 3.\nThe edges in G are: (0, 1) (0, 3) (1, 2)'
' (2, 3).\n'
),
),
dict(
testcase_name='incident_integer',
encoding_method='incident',
expected_result=(
'G describes a graph among nodes 0, 1, 2, and 3.\nIn this'
' graph:\nNode 0 is connected to nodes 1, 3.\nNode 1 is connected'
' to nodes 0, 2.\nNode 2 is connected to nodes 1, 3.\nNode 3 is'
' connected to nodes 2, 0.\n'
),
),
dict(
testcase_name='friendship_per_line_popular',
encoding_method='friendship',
expected_result=(
'G describes a friendship graph among nodes James, Robert, John,'
' and Michael.\nWe have the following edges in G:\nJames and'
' Robert are friends.\nJames and Michael are friends.\nRobert and'
' John are friends.\nJohn and Michael are friends.\n'
),
),
dict(
testcase_name='social_network_politician',
encoding_method='politician',
expected_result=(
'G describes a social network graph among nodes Barack, Jimmy,'
' Arnold, and Bernie.\nWe have the following edges in'
' G:\nBarack and Jimmy are connected.\nBarack and Bernie are'
' connected.\nJimmy and Arnold are connected.\nArnold and'
' Bernie are connected.\n'
),
),
)
def test_encoders(self, encoding_method, expected_result):
self.assertEqual(
graph_text_encoders.encode_graph(_G, encoding_method),
expected_result,
)
if __name__ == '__main__':
googletest.main()
================================================
FILE: talk_like_a_graph/name_dictionaries.py
================================================
"""Creates a dictionary mapping integers to node names."""
import random
_RANDOM_SEED = 1234
random.seed(_RANDOM_SEED)
_INTEGER_NAMES = [str(x) for x in range(10000)]
_POPULAR_NAMES = [
"James",
"Robert",
"John",
"Michael",
"David",
"Mary",
"Patricia",
"Jennifer",
"Linda",
"Elizabeth",
"William",
"Richard",
"Joseph",
"Thomas",
"Christopher",
"Barbara",
"Susan",
"Jessica",
"Sarah",
"Karen",
"Daniel",
"Lisa",
"Matthew",
"Nancy",
"Anthony",
"Betty",
"Mark",
"Margaret",
"Donald",
"Sandra",
"Steven",
"Ashley",
"Paul",
"Kimberly",
"Andrew",
"Emily",
"Joshua",
"Donna",
"Kenneth",
"Michelle",
"Kevin",
"Carol",
"Brian",
"Amanda",
"George",
"Melissa",
"Edward",
"Deborah",
"Ronald",
"Stephanie",
"Timothy",
"Rebecca",
"Jason",
"Sharon",
"Jeffrey",
"Laura",
"Ryan",
"Cynthia",
"Jacob",
"Dorothy",
"Gary",
"Olivia",
"Nicholas",
"Emma",
"Eric",
"Sophia",
"Jonathan",
"Ava",
"Stephen",
"Isabella",
"Scott",
"Mia",
"Justin",
"Abigail",
"Brandon",
"Madison",
"Frank",
"Chloe",
"Benjamin",
"Victoria",
"Samuel",
"Lauren",
"Gregory",
"Hannah",
"Alexander",
"Grace",
"Frank",
"Alexis",
"Raymond",
"Alice",
"Patrick",
"Samantha",
"Jack",
"Natalie",
"Dennis",
"Anna",
"Jerry",
"Taylor",
"Tyler",
"Kayla",
"Henry",
"Hailey",
"Douglas",
"Jasmine",
"Peter",
"Nicole",
"Adam",
"Amy",
"Nathan",
"Christina",
"Zachary",
"Andrea",
"Jose",
"Leah",
"Walter",
"Angelina",
"Harold",
"Valerie",
"Kyle",
"Veronica",
"Ethan",
"Carl",
"Arthur",
"Roger",
"Noah",
]
_SOUTH_PARK_NAMES = [
"Eric",
"Kenny",
"Kyle",
"Stan",
"Tolkien",
"Heidi",
"Bebe",
"Liane",
"Sharon",
"Linda",
"Gerald",
"Veronica",
"Michael",
"Jimbo",
"Herbert",
"Malcolm",
"Gary",
"Steve",
"Chris",
"Wendy",
]
_GOT_NAMES = [
"Ned",
"Cat",
"Daenerys",
"Jon",
"Bran",
"Sansa",
"Arya",
"Cersei",
"Jaime",
"Petyr",
"Robert",
"Jorah",
"Viserys",
"Joffrey",
"Maester",
"Theon",
"Rodrik",
"Lysa",
"Stannis",
"Osha",
]
_POLITICIAN_NAMES = [
"Barack",
"Jimmy",
"Arnold",
"Bernie",
"Bill",
"Kamala",
"Hillary",
"Elizabeth",
"John",
"Ben",
"Joe",
"Alexandria",
"George",
"Nancy",
"Pete",
"Madeleine",
"Elijah",
"Gabrielle",
"Al",
]
_ALPHABET_NAMES = [
"A",
"B",
"C",
"D",
"E",
"F",
"G",
"H",
"I",
"J",
"K",
"L",
"M",
"N",
"O",
"P",
"Q",
"R",
"S",
"T",
"U",
"V",
"W",
"X",
"Y",
"Z",
"AA",
"BB",
"CC",
"DD",
"EE",
"FF",
"GG",
"HH",
"II",
"JJ",
"KK",
"LL",
"MM",
"NN",
"OO",
"PP",
"QQ",
"RR",
"SS",
"TT",
"UU",
"VV",
"WW",
"XX",
"YY",
"ZZ",
]
def create_name_dict(graph, name: str, nnodes: int = 20) -> dict[int, str]:
"""The runner function to map integers to node names.
Args:
graph: the graph to be encoded.
name: name of the approach for mapping.
nnodes: optionally provide nnodes in the graph to be encoded.
Returns:
A dictionary from integers to strings.
"""
if name == "alphabet":
names_list = _ALPHABET_NAMES
elif name == "integer":
names_list = _INTEGER_NAMES
elif name == "random_integer":
names_list = []
for _ in range(nnodes):
names_list.append(str(random.randint(0, 1000000)))
elif name == "popular":
names_list = _POPULAR_NAMES
elif name == "south_park":
names_list = _SOUTH_PARK_NAMES
elif name == "got":
names_list = _GOT_NAMES
elif name == "politician":
names_list = _POLITICIAN_NAMES
elif name == "nx_node_name":
return {x: str(x) for x in graph.nodes()}
else:
raise ValueError(f"Unknown approach: {name}")
name_dict = {}
for ind, value in enumerate(names_list):
name_dict[ind] = value
return name_dict
================================================
FILE: tutorial/KDD-Tutorial-1-Talk-Like-a-Graph.ipynb
================================================
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "EMeYjNCilDE_"
},
"source": [
"This is a noteboook that illustrates how to encode structured data for use in Large Language Models. It is Part 1 of a two part tutorial from **KDD'24**.\n",
"\n",
"If you find this tutorial useful or want to know more, please consider our publication: \n",
"[**Talk like a Graph: Encoding Graphs for Large Language Models**](https://openreview.net/pdf?id=IuXR1CCrSi)\n",
"\n",
"```\n",
"@inproceedings{\n",
"fatemi2024talk,\n",
" title={Talk like a Graph: Encoding Graphs for Large Language Models},\n",
" author={Bahare Fatemi and Jonathan Halcrow and Bryan Perozzi},\n",
" booktitle={The Twelfth International Conference on Learning Representations},\n",
" year={2024},\n",
" url={https://openreview.net/forum?id=IuXR1CCrSi}\n",
"}\n",
"```\n",
"\n",
"\n",
"\n",
"\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5Pt0yWFcnIdy"
},
"source": [
"## Tutorial Part I: Text Encoding of Graph Information\n",
"\n",
"This first notebook focuses on using text serialization methods from **Talk Like a Graph** to perform graph reasoning tasks with off-the-shelf LLMs. This notebook focuses on using static **Gemini** models (which have free API quota), but the techniques (and code) here are generally applicable to any LLM.\n",
"\n",
"\n",
"Notebook Outline:\n",
"\n",
"1. Setup (Install Dependencies, get Google Cloud API key)\n",
"1. Dataset creation\n",
"1. Graph-to-Text conversion\n",
"1. Evaluation\n",
"1. Exercise: *Graph Encoding Challenge*\n",
"1. Exercise: DBLP Dataset\n",
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "SAgPqF6SANGs"
},
"source": [
"# Setup"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "N6cnEIIlXhCj"
},
"source": [
"## Get Google AI Studio Key\n",
"\n",
"Go to [Google AI Studio](https://aistudio.google.com/app/u/1/apikey) and click on the `Create API Key` button. Follow prompts in dialog to setup an API key with a new Google Cloud Project.\n",
"\n",
"Then add the API key to your colab secrets (using the side panel 🔑) as in the picture below:\n",
"\n",
"
FILE: talk_like_a_graph/graph_generators.py
function generate_graphs (line 21) | def generate_graphs(
function remove_graph_data (line 142) | def remove_graph_data(graph: nx.Graph) -> nx.Graph:
function randomize_directions (line 152) | def randomize_directions(graph: nx.Graph) -> nx.DiGraph:
function remove_directions (line 166) | def remove_directions(graph: nx.Graph) -> nx.Graph:
FILE: talk_like_a_graph/graph_generators_runner.py
function write_graphs (line 43) | def write_graphs(graphs: list[nx.Graph], output_dir: str) -> None:
function main (line 56) | def main(argv: Sequence[str]) -> None:
FILE: talk_like_a_graph/graph_generators_test.py
class GraphGenerationTest (line 6) | class GraphGenerationTest(absltest.TestCase, parameterized.TestCase):
method test_number_of_graphs (line 34) | def test_number_of_graphs(self, algorithm, directed, k):
method test_directions (line 110) | def test_directions(self, algorithm, directed):
FILE: talk_like_a_graph/graph_metrics.py
function yes_no_accuracy (line 10) | def yes_no_accuracy(
FILE: talk_like_a_graph/graph_metrics_test.py
class GraphTasksTest (line 5) | class GraphTasksTest(absltest.TestCase):
method test_yes_no_correct (line 7) | def test_yes_no_correct(self):
method test_yes_no_ambiguous (line 22) | def test_yes_no_ambiguous(self):
method test_yes_no_indeterminate (line 37) | def test_yes_no_indeterminate(self):
method test_yes_no_accuracy_raises_on_ambiguous_target (line 52) | def test_yes_no_accuracy_raises_on_ambiguous_target(self):
method test_yes_no_accuracy_raises_on_indeterminate_target (line 59) | def test_yes_no_accuracy_raises_on_indeterminate_target(self):
FILE: talk_like_a_graph/graph_tasks.py
class GraphTask (line 11) | class GraphTask:
method __init__ (line 14) | def __init__(self):
method prepare_examples_dict (line 18) | def prepare_examples_dict(
method create_few_shot_example (line 26) | def create_few_shot_example(
class CycleCheck (line 32) | class CycleCheck(GraphTask):
method __init__ (line 35) | def __init__(self):
method prepare_examples_dict (line 40) | def prepare_examples_dict(
method create_few_shot_example (line 69) | def create_few_shot_example(
method choose_few_shot_examples (line 98) | def choose_few_shot_examples(
class EdgeExistence (line 119) | class EdgeExistence(GraphTask):
method __init__ (line 122) | def __init__(self):
method prepare_examples_dict (line 126) | def prepare_examples_dict(
method create_few_shot_example (line 161) | def create_few_shot_example(
class NodeCount (line 192) | class NodeCount(GraphTask):
method __init__ (line 195) | def __init__(self):
method prepare_examples_dict (line 200) | def prepare_examples_dict(
method get_nodes_string (line 223) | def get_nodes_string(self, name_dict: dict[int, str], nnodes: int) -> ...
method create_few_shot_example (line 230) | def create_few_shot_example(
class NodeDegree (line 247) | class NodeDegree(GraphTask):
method __init__ (line 250) | def __init__(self):
method prepare_examples_dict (line 254) | def prepare_examples_dict(
method get_edge_string (line 282) | def get_edge_string(
method create_few_shot_example (line 299) | def create_few_shot_example(
class EdgeCount (line 319) | class EdgeCount(GraphTask):
method __init__ (line 322) | def __init__(self):
method prepare_examples_dict (line 327) | def prepare_examples_dict(
method get_edges_string (line 350) | def get_edges_string(
method create_few_shot_example (line 360) | def create_few_shot_example(
class ConnectedNodes (line 376) | class ConnectedNodes(GraphTask):
method __init__ (line 379) | def __init__(self):
method prepare_examples_dict (line 383) | def prepare_examples_dict(
method get_connected_nodes (line 416) | def get_connected_nodes(
method create_few_shot_example (line 433) | def create_few_shot_example(
class DisconnectedNodes (line 464) | class DisconnectedNodes(GraphTask):
method __init__ (line 467) | def __init__(self):
method prepare_examples_dict (line 471) | def prepare_examples_dict(
method get_disconnected_nodes (line 508) | def get_disconnected_nodes(
method create_few_shot_example (line 537) | def create_few_shot_example(
class Reachability (line 574) | class Reachability(GraphTask):
method __init__ (line 577) | def __init__(self):
method prepare_examples_dict (line 581) | def prepare_examples_dict(
method create_few_shot_example (line 614) | def create_few_shot_example(
class ShortestPath (line 655) | class ShortestPath(GraphTask):
method __init__ (line 658) | def __init__(self):
method prepare_examples_dict (line 662) | def prepare_examples_dict(
method create_few_shot_example (line 703) | def create_few_shot_example(
class TriangleCounting (line 751) | class TriangleCounting(GraphTask):
method __init__ (line 754) | def __init__(self):
method prepare_examples_dict (line 759) | def prepare_examples_dict(
method create_few_shot_example (line 786) | def create_few_shot_example(
class MaximumFlow (line 829) | class MaximumFlow(GraphTask):
method __init__ (line 832) | def __init__(self):
method prepare_examples_dict (line 836) | def prepare_examples_dict(
method create_few_shot_example (line 870) | def create_few_shot_example(
function has_edge_weights (line 912) | def has_edge_weights(graph):
function add_edge_weight (line 919) | def add_edge_weight(graph):
class NodeClassification (line 928) | class NodeClassification(GraphTask):
method __init__ (line 931) | def __init__(self):
method prepare_examples_dict (line 943) | def prepare_examples_dict(
method create_few_shot_example (line 988) | def create_few_shot_example(
FILE: talk_like_a_graph/graph_tasks_generator.py
function zero_shot (line 35) | def zero_shot(
function few_shot (line 69) | def few_shot(
function generate_random_sbm_graph (line 116) | def generate_random_sbm_graph(random_state: np.random.RandomState):
function main (line 129) | def main(argv: Sequence[str]) -> None:
FILE: talk_like_a_graph/graph_tasks_utils.py
function laplacian_pos_embedding (line 19) | def laplacian_pos_embedding(graph: nx.Graph, units: int = 4) -> nx.Graph:
function to_tfgnn (line 33) | def to_tfgnn(graph: nx.Graph, node_ids: list[int]) -> tfgnn.GraphTensor:
function create_example_feature (line 116) | def create_example_feature(
function load_graphs (line 179) | def load_graphs(
function prepare_examples (line 204) | def prepare_examples(
function create_zero_shot_task (line 247) | def create_zero_shot_task(
function write_examples (line 267) | def write_examples(examples: list[example_pb2.Example], output_path: str):
function prepare_few_shots (line 273) | def prepare_few_shots(
function choose_few_shot_examples (line 291) | def choose_few_shot_examples(
function create_few_shot_task (line 304) | def create_few_shot_task(
FILE: talk_like_a_graph/graph_text_encoders.py
function create_node_string (line 8) | def create_node_string(name_dict, nnodes: int) -> str:
function nx_encoder (line 17) | def nx_encoder(graph: nx.Graph, _: dict[int, str], edge_type="id") -> str:
function adjacency_encoder (line 44) | def adjacency_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
function friendship_encoder (line 65) | def friendship_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
function coauthorship_encoder (line 80) | def coauthorship_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> ...
function incident_encoder (line 99) | def incident_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
function social_network_encoder (line 125) | def social_network_encoder(graph: nx.Graph, name_dict: dict[int, str]) -...
function expert_encoder (line 143) | def expert_encoder(graph: nx.Graph, name_dict: dict[int, str]) -> str:
function nodes_to_text (line 156) | def nodes_to_text(graph, encoding_type):
function get_tlag_node_encoder (line 180) | def get_tlag_node_encoder(graph, encoder_name):
function with_ids (line 224) | def with_ids(graph: nx.Graph, node_encoder: str) -> nx.Graph:
function encode_graph (line 229) | def encode_graph(
FILE: talk_like_a_graph/graph_text_encoders_test.py
class GraphTextEncodersTest (line 20) | class GraphTextEncodersTest(absltest.TestCase, parameterized.TestCase):
method test_encoders (line 65) | def test_encoders(self, encoding_method, expected_result):
FILE: talk_like_a_graph/name_dictionaries.py
function create_name_dict (line 265) | def create_name_dict(graph, name: str, nnodes: int = 20) -> dict[int, str]:
Condensed preview — 20 files, each showing path, character count, and a content snippet. Download the .json file or copy for the full structured content (362K chars).
[
{
"path": "CONTRIBUTING.md",
"chars": 969,
"preview": "# How to Contribute\n\n## Contributor License Agreement\n\nContributions to this project must be accompanied by a Contributo"
},
{
"path": "LICENSE",
"chars": 11357,
"preview": " Apache License\n Version 2.0, January 2004\n "
},
{
"path": "README.md",
"chars": 1223,
"preview": "# Encoding Graphs and Structured Data in Language Models\n\nThis repository contains code for\n[Talk like a Graph: Encoding"
},
{
"path": "__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "talk_like_a_graph/README.md",
"chars": 1306,
"preview": "# Using Large Language Models to Solve Graph Problems\n\nThis repository contains the code to generate graph reasoning pro"
},
{
"path": "talk_like_a_graph/__init__.py",
"chars": 0,
"preview": ""
},
{
"path": "talk_like_a_graph/graph_generators.py",
"chars": 5972,
"preview": "r\"\"\"Random graph generation.\"\"\"\n\nimport random\n\nimport networkx as nx\nimport numpy as np\n\n\n_NUMBER_OF_NODES_RANGE = {\n "
},
{
"path": "talk_like_a_graph/graph_generators_runner.py",
"chars": 2266,
"preview": "r\"\"\"Random graph generation.\n\nThis code generates random graph using different algorithms.\n\n# Placeholder for Google-int"
},
{
"path": "talk_like_a_graph/graph_generators_test.py",
"chars": 2817,
"preview": "from absl.testing import parameterized\nfrom . import graph_generators\nfrom absl.testing import absltest\n\n\nclass GraphGen"
},
{
"path": "talk_like_a_graph/graph_metrics.py",
"chars": 2263,
"preview": "\"\"\"Metrics for seqio tasks over graph data.\n\nThis module contains definitions of metric_fns to be used for scoring\ngraph"
},
{
"path": "talk_like_a_graph/graph_metrics_test.py",
"chars": 2056,
"preview": "from . import graph_metrics\nfrom absl.testing import absltest\n\n\nclass GraphTasksTest(absltest.TestCase):\n\n def test_yes"
},
{
"path": "talk_like_a_graph/graph_tasks.py",
"chars": 33464,
"preview": "\"\"\"The graph tasks to be tried with LLMs.\"\"\"\n\nimport random\n\nimport networkx as nx\nimport numpy as np\n\nfrom . import gra"
},
{
"path": "talk_like_a_graph/graph_tasks_generator.py",
"chars": 6385,
"preview": "r\"\"\"The graph tasks to be tried with LLMs..\n\nThis code loads graphs and creates graph tasks and output them as tf exampl"
},
{
"path": "talk_like_a_graph/graph_tasks_utils.py",
"chars": 11002,
"preview": "\"\"\"The graph tasks to be tried with LLMs.\"\"\"\n\nimport os\nimport random\n\nimport networkx as nx\nimport numpy as np\nimport s"
},
{
"path": "talk_like_a_graph/graph_text_encoders.py",
"chars": 10727,
"preview": "\"\"\"Library for encoding graphs in text.\"\"\"\n\nimport networkx as nx\n\nfrom . import name_dictionaries\n\n\ndef create_node_str"
},
{
"path": "talk_like_a_graph/graph_text_encoders_test.py",
"chars": 2448,
"preview": "\"\"\"Testing for graph_text_encoders.py.\"\"\"\n\nfrom absl.testing import parameterized\nimport networkx as nx\n\nfrom . import g"
},
{
"path": "talk_like_a_graph/name_dictionaries.py",
"chars": 4404,
"preview": "\"\"\"Creates a dictionary mapping integers to node names.\"\"\"\n\nimport random\n\n_RANDOM_SEED = 1234\nrandom.seed(_RANDOM_SEED)"
},
{
"path": "tutorial/KDD-Tutorial-1-Talk-Like-a-Graph.ipynb",
"chars": 198948,
"preview": "{\n \"cells\": [\n {\n \"cell_type\": \"markdown\",\n \"metadata\": {\n \"id\": \"EMeYjNCilDE_\"\n },\n \"sou"
},
{
"path": "tutorial/KDD-Tutorial-2-Let-Your-Graph-Do-The-Talking.ipynb",
"chars": 48987,
"preview": "{\n \"nbformat\": 4,\n \"nbformat_minor\": 0,\n \"metadata\": {\n \"colab\": {\n \"provenance\": [],\n \"gpuType\": \"V28\"\n"
},
{
"path": "tutorial/README.md",
"chars": 1907,
"preview": "# <p align=\"center\"> Tutorial on Graph Reasoning with LLMs (GReaL) </p>\n\n## <p align=\"center\">**KDD'24 Tutorial**</p>\n\n#"
}
]
About this extraction
This page contains the full source code of the google-research/talk-like-a-graph GitHub repository, extracted and formatted as plain text for AI agents and large language models (LLMs). The extraction includes 20 files (340.3 KB), approximately 158.0k tokens, and a symbol index with 105 extracted functions, classes, methods, constants, and types. Use this with OpenClaw, Claude, ChatGPT, Cursor, Windsurf, or any other AI tool that accepts text input. You can copy the full output to your clipboard or download it as a .txt file.
Extracted by GitExtract — free GitHub repo to text converter for AI. Built by Nikandr Surkov.