wehub-resource-sync
61 lines
1.6 KiB
Python
61 lines
1.6 KiB
Python
""" load dataset from ogb """
|
|
|
|
import argparse
|
|
import time
|
|
|
|
from ogb.linkproppred import DglLinkPropPredDataset
|
|
|
|
|
|
def load_from_ogbl_with_name(name):
|
|
choices = ["ogbl-collab", "ogbl-ddi", "ogbl-ppa", "ogbl-citation"]
|
|
assert name in choices, "name must be selected from " + str(choices)
|
|
dataset = DglLinkPropPredDataset(name)
|
|
return dataset[0]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"--name",
|
|
type=str,
|
|
choices=["ogbl-collab", "ogbl-ddi", "ogbl-ppa", "ogbl-citation"],
|
|
default="ogbl-collab",
|
|
help="name of datasets by ogb",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
print("loading graph... it might take some time")
|
|
name = args.name
|
|
g = load_from_ogbl_with_name(name=name)
|
|
|
|
try:
|
|
w = g.edata["edge_weight"]
|
|
weighted = True
|
|
except:
|
|
weighted = False
|
|
|
|
edge_num = g.edges()[0].shape[0]
|
|
src = list(g.edges()[0])
|
|
tgt = list(g.edges()[1])
|
|
if weighted:
|
|
weight = list(g.edata["edge_weight"])
|
|
|
|
print("writing...")
|
|
start_time = time.time()
|
|
with open(name + "-net.txt", "w") as f:
|
|
for i in range(edge_num):
|
|
if weighted:
|
|
f.write(
|
|
str(src[i].item())
|
|
+ " "
|
|
+ str(tgt[i].item())
|
|
+ " "
|
|
+ str(weight[i].item())
|
|
+ "\n"
|
|
)
|
|
else:
|
|
f.write(
|
|
str(src[i].item()) + " " + str(tgt[i].item()) + " " + "1\n"
|
|
)
|
|
print("writing used time: %d s" % int(time.time() - start_time))
|