Create Freebase ID to node index mapping for FB15k-237 dataset.

fb15k_mapping.py

"""Create Freebase ID to node index mapping for FB15k-237 dataset."""

import json
from pathlib import Path
from typing import Dict

from torch_geometric.datasets import FB15k_237

NODE_PATH = Path("nodes.json")
EDGE_PATH = Path("edges.json")

ds = FB15k_237(root="data")

node_dict: Dict[str, int] = {}
rel_dict: Dict[str, int] = {}

for path in ds.raw_paths:
    with open(path, "r") as f:
        lines = [x.split("\t") for x in f.read().split("\n")[:-1]]

    for i, (src, rel, dst) in enumerate(lines):
        if src not in node_dict:
            node_dict[src] = len(node_dict)
        if dst not in node_dict:
            node_dict[dst] = len(node_dict)
        if rel not in rel_dict:
            rel_dict[rel] = len(rel_dict)


with open(NODE_PATH, "w") as f:
    json.dump(node_dict, f)

with open(EDGE_PATH, "w") as f:
    json.dump(rel_dict, f)