Source code for dgl.data.rdf

"""RDF datasets
Datasets from "A Collection of Benchmark Datasets for
Systematic Evaluations of Machine Learning on
the Semantic Web"
"""
import abc
import itertools
import os
import re
from collections import OrderedDict

import networkx as nx
import numpy as np

import dgl
import dgl.backend as F

from .dgl_dataset import DGLBuiltinDataset
from .utils import (
    _get_dgl_url,
    generate_mask_tensor,
    idx2mask,
    load_graphs,
    load_info,
    save_graphs,
    save_info,
)

__all__ = ["AIFBDataset", "MUTAGDataset", "BGSDataset", "AMDataset"]

# Dictionary for renaming reserved node/edge type names to the ones
# that are allowed by nn.Module.
RENAME_DICT = {
    "type": "rdftype",
    "rev-type": "rev-rdftype",
}


class Entity:
    """Class for entities
    Parameters
    ----------
    id : str
        ID of this entity
    cls : str
        Type of this entity
    """

    def __init__(self, e_id, cls):
        self.id = e_id
        self.cls = cls

    def __str__(self):
        return "{}/{}".format(self.cls, self.id)


class Relation:
    """Class for relations
    Parameters
    ----------
    cls : str
        Type of this relation
    """

    def __init__(self, cls):
        self.cls = cls

    def __str__(self):
        return str(self.cls)


class RDFGraphDataset(DGLBuiltinDataset):
    """Base graph dataset class from RDF tuples.

    To derive from this, implement the following abstract methods:
    * ``parse_entity``
    * ``parse_relation``
    * ``process_tuple``
    * ``process_idx_file_line``
    * ``predict_category``
    Preprocessed graph and other data will be cached in the download folder
    to speedup data loading.
    The dataset should contain a "trainingSet.tsv" and a "testSet.tsv" file
    for training and testing samples.

    Attributes
    ----------
    num_classes : int
        Number of classes to predict
    predict_category : str
        The entity category (node type) that has labels for prediction

    Parameters
    ----------
    name : str
        Name of the dataset
    url : str or path
        URL to download the raw dataset.
    predict_category : str
        Predict category.
    print_every : int, optional
        Preprocessing log for every X tuples.
    insert_reverse : bool, optional
        If true, add reverse edge and reverse relations to the final graph.
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool, optional
        If true, force load and process from raw data. Ignore cached pre-processed data.
    verbose : bool
        Whether to print out progress information. Default: True.
    transform : callable, optional
        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
        a transformed version. The :class:`~dgl.DGLGraph` object will be
        transformed before every access.
    """

    def __init__(
        self,
        name,
        url,
        predict_category,
        print_every=10000,
        insert_reverse=True,
        raw_dir=None,
        force_reload=False,
        verbose=True,
        transform=None,
    ):
        self._insert_reverse = insert_reverse
        self._print_every = print_every
        self._predict_category = predict_category

        super(RDFGraphDataset, self).__init__(
            name,
            url,
            raw_dir=raw_dir,
            force_reload=force_reload,
            verbose=verbose,
            transform=transform,
        )

    def process(self):
        raw_tuples = self.load_raw_tuples(self.raw_path)
        self.process_raw_tuples(raw_tuples, self.raw_path)

    def load_raw_tuples(self, root_path):
        """Loading raw RDF dataset

        Parameters
        ----------
        root_path : str
            Root path containing the data

        Returns
        -------
            Loaded rdf data
        """
        import rdflib as rdf

        raw_rdf_graphs = []
        for _, filename in enumerate(os.listdir(root_path)):
            fmt = None
            if filename.endswith("nt"):
                fmt = "nt"
            elif filename.endswith("n3"):
                fmt = "n3"
            if fmt is None:
                continue
            g = rdf.Graph()
            print("Parsing file %s ..." % filename)
            g.parse(os.path.join(root_path, filename), format=fmt)
            raw_rdf_graphs.append(g)
        return itertools.chain(*raw_rdf_graphs)

    def process_raw_tuples(self, raw_tuples, root_path):
        """Processing raw RDF dataset

        Parameters
        ----------
        raw_tuples:
            Raw rdf tuples
        root_path: str
            Root path containing the data
        """
        mg = nx.MultiDiGraph()
        ent_classes = OrderedDict()
        rel_classes = OrderedDict()
        entities = OrderedDict()
        src = []
        dst = []
        ntid = []
        etid = []
        sorted_tuples = []
        for t in raw_tuples:
            sorted_tuples.append(t)
        sorted_tuples.sort()

        for i, (sbj, pred, obj) in enumerate(sorted_tuples):
            if self.verbose and i % self._print_every == 0:
                print(
                    "Processed %d tuples, found %d valid tuples."
                    % (i, len(src))
                )
            sbjent = self.parse_entity(sbj)
            rel = self.parse_relation(pred)
            objent = self.parse_entity(obj)
            processed = self.process_tuple(
                (sbj, pred, obj), sbjent, rel, objent
            )
            if processed is None:
                # ignored
                continue
            # meta graph
            sbjclsid = _get_id(ent_classes, sbjent.cls)
            objclsid = _get_id(ent_classes, objent.cls)
            relclsid = _get_id(rel_classes, rel.cls)
            mg.add_edge(sbjent.cls, objent.cls, key=rel.cls)
            if self._insert_reverse:
                mg.add_edge(objent.cls, sbjent.cls, key="rev-%s" % rel.cls)
            # instance graph
            src_id = _get_id(entities, str(sbjent))
            if len(entities) > len(ntid):  # found new entity
                ntid.append(sbjclsid)
            dst_id = _get_id(entities, str(objent))
            if len(entities) > len(ntid):  # found new entity
                ntid.append(objclsid)
            src.append(src_id)
            dst.append(dst_id)
            etid.append(relclsid)

        src = np.asarray(src)
        dst = np.asarray(dst)
        ntid = np.asarray(ntid)
        etid = np.asarray(etid)
        ntypes = list(ent_classes.keys())
        etypes = list(rel_classes.keys())

        # add reverse edge with reverse relation
        if self._insert_reverse:
            if self.verbose:
                print("Adding reverse edges ...")
            newsrc = np.hstack([src, dst])
            newdst = np.hstack([dst, src])
            src = newsrc
            dst = newdst
            etid = np.hstack([etid, etid + len(etypes)])
            etypes.extend(["rev-%s" % t for t in etypes])

        hg = self.build_graph(mg, src, dst, ntid, etid, ntypes, etypes)

        if self.verbose:
            print("Load training/validation/testing split ...")
        idmap = F.asnumpy(hg.nodes[self.predict_category].data[dgl.NID])
        glb2lcl = {glbid: lclid for lclid, glbid in enumerate(idmap)}

        def findidfn(ent):
            if ent not in entities:
                return None
            else:
                return glb2lcl[entities[ent]]

        self._hg = hg
        train_idx, test_idx, labels, num_classes = self.load_data_split(
            findidfn, root_path
        )

        train_mask = idx2mask(
            train_idx, self._hg.num_nodes(self.predict_category)
        )
        test_mask = idx2mask(
            test_idx, self._hg.num_nodes(self.predict_category)
        )
        labels = F.tensor(labels, F.data_type_dict["int64"])

        train_mask = generate_mask_tensor(train_mask)
        test_mask = generate_mask_tensor(test_mask)
        self._hg.nodes[self.predict_category].data["train_mask"] = train_mask
        self._hg.nodes[self.predict_category].data["test_mask"] = test_mask
        # TODO(minjie): Deprecate 'labels', use 'label' for consistency.
        self._hg.nodes[self.predict_category].data["labels"] = labels
        self._hg.nodes[self.predict_category].data["label"] = labels
        self._num_classes = num_classes

    def build_graph(self, mg, src, dst, ntid, etid, ntypes, etypes):
        """Build the graphs

        Parameters
        ----------
        mg: MultiDiGraph
            Input graph
        src: Numpy array
            Source nodes
        dst: Numpy array
            Destination nodes
        ntid: Numpy array
            Node types for each node
        etid: Numpy array
            Edge types for each edge
        ntypes: list
            Node types
        etypes: list
            Edge types

        Returns
        -------
        g: DGLGraph
        """
        # create homo graph
        if self.verbose:
            print("Creating one whole graph ...")
        g = dgl.graph((src, dst))
        g.ndata[dgl.NTYPE] = F.tensor(ntid)
        g.edata[dgl.ETYPE] = F.tensor(etid)
        if self.verbose:
            print("Total #nodes:", g.num_nodes())
            print("Total #edges:", g.num_edges())

        # rename names such as 'type' so that they an be used as keys
        # to nn.ModuleDict
        etypes = [RENAME_DICT.get(ty, ty) for ty in etypes]
        mg_edges = mg.edges(keys=True)
        mg = nx.MultiDiGraph()
        for sty, dty, ety in mg_edges:
            mg.add_edge(sty, dty, key=RENAME_DICT.get(ety, ety))

        # convert to heterograph
        if self.verbose:
            print("Convert to heterograph ...")
        hg = dgl.to_heterogeneous(g, ntypes, etypes, metagraph=mg)
        if self.verbose:
            print("#Node types:", len(hg.ntypes))
            print("#Canonical edge types:", len(hg.etypes))
            print("#Unique edge type names:", len(set(hg.etypes)))
        return hg

    def load_data_split(self, ent2id, root_path):
        """Load data split

        Parameters
        ----------
        ent2id: func
            A function mapping entity to id
        root_path: str
            Root path containing the data

        Return
        ------
        train_idx: Numpy array
            Training set
        test_idx: Numpy array
            Testing set
        labels: Numpy array
            Labels
        num_classes: int
            Number of classes
        """
        label_dict = {}
        labels = np.zeros((self._hg.num_nodes(self.predict_category),)) - 1
        train_idx = self.parse_idx_file(
            os.path.join(root_path, "trainingSet.tsv"),
            ent2id,
            label_dict,
            labels,
        )
        test_idx = self.parse_idx_file(
            os.path.join(root_path, "testSet.tsv"), ent2id, label_dict, labels
        )
        train_idx = np.array(train_idx)
        test_idx = np.array(test_idx)
        labels = np.array(labels)
        num_classes = len(label_dict)
        return train_idx, test_idx, labels, num_classes

    def parse_idx_file(self, filename, ent2id, label_dict, labels):
        """Parse idx files

        Parameters
        ----------
        filename: str
            File to parse
        ent2id: func
            A function mapping entity to id
        label_dict: dict
            Map label to label id
        labels: dict
            Map entity id to label id

        Return
        ------
        idx: list
            Entity idss
        """
        idx = []
        with open(filename, "r") as f:
            for i, line in enumerate(f):
                if i == 0:
                    continue  # first line is the header
                sample, label = self.process_idx_file_line(line)
                # person, _, label = line.strip().split('\t')
                ent = self.parse_entity(sample)
                entid = ent2id(str(ent))
                if entid is None:
                    print(
                        'Warning: entity "%s" does not have any valid links associated. Ignored.'
                        % str(ent)
                    )
                else:
                    idx.append(entid)
                    lblid = _get_id(label_dict, label)
                    labels[entid] = lblid
        return idx

    def has_cache(self):
        """check if there is a processed data"""
        graph_path = os.path.join(self.save_path, self.save_name + ".bin")
        info_path = os.path.join(self.save_path, self.save_name + ".pkl")
        if os.path.exists(graph_path) and os.path.exists(info_path):
            return True

        return False

    def save(self):
        """save the graph list and the labels"""
        graph_path = os.path.join(self.save_path, self.save_name + ".bin")
        info_path = os.path.join(self.save_path, self.save_name + ".pkl")
        save_graphs(str(graph_path), self._hg)
        save_info(
            str(info_path),
            {
                "num_classes": self.num_classes,
                "predict_category": self.predict_category,
            },
        )

    def load(self):
        """load the graph list and the labels from disk"""
        graph_path = os.path.join(self.save_path, self.save_name + ".bin")
        info_path = os.path.join(self.save_path, self.save_name + ".pkl")
        graphs, _ = load_graphs(str(graph_path))

        info = load_info(str(info_path))
        self._num_classes = info["num_classes"]
        self._predict_category = info["predict_category"]
        self._hg = graphs[0]
        # For backward compatibility
        if "label" not in self._hg.nodes[self.predict_category].data:
            self._hg.nodes[self.predict_category].data[
                "label"
            ] = self._hg.nodes[self.predict_category].data["labels"]

    def __getitem__(self, idx):
        r"""Gets the graph object"""
        g = self._hg
        if self._transform is not None:
            g = self._transform(g)
        return g

    def __len__(self):
        r"""The number of graphs in the dataset."""
        return 1

    @property
    def save_name(self):
        return self.name + "_dgl_graph"

    @property
    def predict_category(self):
        return self._predict_category

    @property
    def num_classes(self):
        return self._num_classes

    @abc.abstractmethod
    def parse_entity(self, term):
        """Parse one entity from an RDF term.
        Return None if the term does not represent a valid entity and the
        whole tuple should be ignored.
        Parameters
        ----------
        term : rdflib.term.Identifier
            RDF term
        Returns
        -------
        Entity or None
            An entity.
        """
        pass

    @abc.abstractmethod
    def parse_relation(self, term):
        """Parse one relation from an RDF term.
        Return None if the term does not represent a valid relation and the
        whole tuple should be ignored.
        Parameters
        ----------
        term : rdflib.term.Identifier
            RDF term
        Returns
        -------
        Relation or None
            A relation
        """
        pass

    @abc.abstractmethod
    def process_tuple(self, raw_tuple, sbj, rel, obj):
        """Process the tuple.
        Return (Entity, Relation, Entity) tuple for as the final tuple.
        Return None if the tuple should be ignored.

        Parameters
        ----------
        raw_tuple : tuple of rdflib.term.Identifier
            (subject, predicate, object) tuple
        sbj : Entity
            Subject entity
        rel : Relation
            Relation
        obj : Entity
            Object entity
        Returns
        -------
        (Entity, Relation, Entity)
            The final tuple or None if should be ignored
        """
        pass

    @abc.abstractmethod
    def process_idx_file_line(self, line):
        """Process one line of ``trainingSet.tsv`` or ``testSet.tsv``.
        Parameters
        ----------
        line : str
            One line of the file
        Returns
        -------
        (str, str)
            One sample and its label
        """
        pass


def _get_id(dict, key):
    id = dict.get(key, None)
    if id is None:
        id = len(dict)
        dict[key] = id
    return id




class AIFBDataset(RDFGraphDataset):
    r"""AIFB dataset for node classification task

    AIFB DataSet is a Semantic Web (RDF) dataset used as a benchmark in
    data mining.  It records the organizational structure of AIFB at the
    University of Karlsruhe.

    AIFB dataset statistics:

    - Nodes: 7262
    - Edges: 48810 (including reverse edges)
    - Target Category: Personen
    - Number of Classes: 4
    - Label Split:

        - Train: 140
        - Test: 36

    Parameters
    -----------
    print_every : int
        Preprocessing log for every X tuples. Default: 10000.
    insert_reverse : bool
        If true, add reverse edge and reverse relations to the final graph. Default: True.
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool
        Whether to reload the dataset. Default: False
    verbose : bool
        Whether to print out progress information. Default: True.
    transform : callable, optional
        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
        a transformed version. The :class:`~dgl.DGLGraph` object will be
        transformed before every access.

    Attributes
    ----------
    num_classes : int
        Number of classes to predict
    predict_category : str
        The entity category (node type) that has labels for prediction

    Examples
    --------
    >>> dataset = dgl.data.rdf.AIFBDataset()
    >>> graph = dataset[0]
    >>> category = dataset.predict_category
    >>> num_classes = dataset.num_classes
    >>>
    >>> train_mask = g.nodes[category].data['train_mask']
    >>> test_mask = g.nodes[category].data['test_mask']
    >>> label = g.nodes[category].data['label']
    """

    entity_prefix = "http://www.aifb.uni-karlsruhe.de/"
    relation_prefix = "http://swrc.ontoware.org/"

    def __init__(
        self,
        print_every=10000,
        insert_reverse=True,
        raw_dir=None,
        force_reload=False,
        verbose=True,
        transform=None,
    ):
        import rdflib as rdf

        self.employs = rdf.term.URIRef(
            "http://swrc.ontoware.org/ontology#employs"
        )
        self.affiliation = rdf.term.URIRef(
            "http://swrc.ontoware.org/ontology#affiliation"
        )
        url = _get_dgl_url("dataset/rdf/aifb-hetero.zip")
        name = "aifb-hetero"
        predict_category = "Personen"
        super(AIFBDataset, self).__init__(
            name,
            url,
            predict_category,
            print_every=print_every,
            insert_reverse=insert_reverse,
            raw_dir=raw_dir,
            force_reload=force_reload,
            verbose=verbose,
            transform=transform,
        )



    def __getitem__(self, idx):
        r"""Gets the graph object

        Parameters
        -----------
        idx: int
            Item index, AIFBDataset has only one graph object

        Return
        -------
        :class:`dgl.DGLGraph`

            The graph contains:

            - ``ndata['train_mask']``: mask for training node set
            - ``ndata['test_mask']``: mask for testing node set
            - ``ndata['label']``: node labels
        """
        return super(AIFBDataset, self).__getitem__(idx)




    def __len__(self):
        r"""The number of graphs in the dataset.

        Return
        -------
        int
        """
        return super(AIFBDataset, self).__len__()


    def parse_entity(self, term):
        import rdflib as rdf

        if isinstance(term, rdf.Literal):
            return Entity(e_id=str(term), cls="_Literal")
        if isinstance(term, rdf.BNode):
            return None
        entstr = str(term)
        if entstr.startswith(self.entity_prefix):
            sp = entstr.split("/")
            return Entity(e_id=sp[5], cls=sp[3])
        else:
            return None

    def parse_relation(self, term):
        if term == self.employs or term == self.affiliation:
            return None
        relstr = str(term)
        if relstr.startswith(self.relation_prefix):
            return Relation(cls=relstr.split("/")[3])
        else:
            relstr = relstr.split("/")[-1]
            return Relation(cls=relstr)

    def process_tuple(self, raw_tuple, sbj, rel, obj):
        if sbj is None or rel is None or obj is None:
            return None
        return (sbj, rel, obj)

    def process_idx_file_line(self, line):
        person, _, label = line.strip().split("\t")
        return person, label





class MUTAGDataset(RDFGraphDataset):
    r"""MUTAG dataset for node classification task

    Mutag dataset statistics:

    - Nodes: 27163
    - Edges: 148100 (including reverse edges)
    - Target Category: d
    - Number of Classes: 2
    - Label Split:

        - Train: 272
        - Test: 68

    Parameters
    -----------
    print_every : int
        Preprocessing log for every X tuples. Default: 10000.
    insert_reverse : bool
        If true, add reverse edge and reverse relations to the final graph. Default: True.
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool
        Whether to reload the dataset. Default: False
    verbose : bool
        Whether to print out progress information. Default: True.
    transform : callable, optional
        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
        a transformed version. The :class:`~dgl.DGLGraph` object will be
        transformed before every access.

    Attributes
    ----------
    num_classes : int
        Number of classes to predict
    predict_category : str
        The entity category (node type) that has labels for prediction
    graph : :class:`dgl.DGLGraph`
        Graph structure

    Examples
    --------
    >>> dataset = dgl.data.rdf.MUTAGDataset()
    >>> graph = dataset[0]
    >>> category = dataset.predict_category
    >>> num_classes = dataset.num_classes
    >>>
    >>> train_mask = g.nodes[category].data['train_mask']
    >>> test_mask = g.nodes[category].data['test_mask']
    >>> label = g.nodes[category].data['label']
    """

    d_entity = re.compile("d[0-9]")
    bond_entity = re.compile("bond[0-9]")

    entity_prefix = "http://dl-learner.org/carcinogenesis#"
    relation_prefix = entity_prefix

    def __init__(
        self,
        print_every=10000,
        insert_reverse=True,
        raw_dir=None,
        force_reload=False,
        verbose=True,
        transform=None,
    ):
        import rdflib as rdf

        self.is_mutagenic = rdf.term.URIRef(
            "http://dl-learner.org/carcinogenesis#isMutagenic"
        )
        self.rdf_type = rdf.term.URIRef(
            "http://www.w3.org/1999/02/22-rdf-syntax-ns#type"
        )
        self.rdf_subclassof = rdf.term.URIRef(
            "http://www.w3.org/2000/01/rdf-schema#subClassOf"
        )
        self.rdf_domain = rdf.term.URIRef(
            "http://www.w3.org/2000/01/rdf-schema#domain"
        )

        url = _get_dgl_url("dataset/rdf/mutag-hetero.zip")
        name = "mutag-hetero"
        predict_category = "d"
        super(MUTAGDataset, self).__init__(
            name,
            url,
            predict_category,
            print_every=print_every,
            insert_reverse=insert_reverse,
            raw_dir=raw_dir,
            force_reload=force_reload,
            verbose=verbose,
            transform=transform,
        )



    def __getitem__(self, idx):
        r"""Gets the graph object

        Parameters
        -----------
        idx: int
            Item index, MUTAGDataset has only one graph object

        Return
        -------
        :class:`dgl.DGLGraph`

            The graph contains:

            - ``ndata['train_mask']``: mask for training node set
            - ``ndata['test_mask']``: mask for testing node set
            - ``ndata['label']``: node labels
        """
        return super(MUTAGDataset, self).__getitem__(idx)




    def __len__(self):
        r"""The number of graphs in the dataset.

        Return
        -------
        int
        """
        return super(MUTAGDataset, self).__len__()


    def parse_entity(self, term):
        import rdflib as rdf

        if isinstance(term, rdf.Literal):
            return Entity(e_id=str(term), cls="_Literal")
        elif isinstance(term, rdf.BNode):
            return None
        entstr = str(term)
        if entstr.startswith(self.entity_prefix):
            inst = entstr[len(self.entity_prefix) :]
            if self.d_entity.match(inst):
                cls = "d"
            elif self.bond_entity.match(inst):
                cls = "bond"
            else:
                cls = None
            return Entity(e_id=inst, cls=cls)
        else:
            return None

    def parse_relation(self, term):
        if term == self.is_mutagenic:
            return None
        relstr = str(term)
        if relstr.startswith(self.relation_prefix):
            cls = relstr[len(self.relation_prefix) :]
            return Relation(cls=cls)
        else:
            relstr = relstr.split("/")[-1]
            return Relation(cls=relstr)

    def process_tuple(self, raw_tuple, sbj, rel, obj):
        if sbj is None or rel is None or obj is None:
            return None

        if not raw_tuple[1].startswith("http://dl-learner.org/carcinogenesis#"):
            obj.cls = "SCHEMA"
            if sbj.cls is None:
                sbj.cls = "SCHEMA"
        if obj.cls is None:
            obj.cls = rel.cls

        assert sbj.cls is not None and obj.cls is not None

        return (sbj, rel, obj)

    def process_idx_file_line(self, line):
        bond, _, label = line.strip().split("\t")
        return bond, label





class BGSDataset(RDFGraphDataset):
    r"""BGS dataset for node classification task

    BGS namespace convention:
    ``http://data.bgs.ac.uk/(ref|id)/<Major Concept>/<Sub Concept>/INSTANCE``.
    We ignored all literal nodes and the relations connecting them in the
    output graph. We also ignored the relation used to mark whether a
    term is CURRENT or DEPRECATED.

    BGS dataset statistics:

    - Nodes: 94806
    - Edges: 672884 (including reverse edges)
    - Target Category: Lexicon/NamedRockUnit
    - Number of Classes: 2
    - Label Split:

        - Train: 117
        - Test: 29

    Parameters
    -----------
    print_every : int
        Preprocessing log for every X tuples. Default: 10000.
    insert_reverse : bool
        If true, add reverse edge and reverse relations to the final graph. Default: True.
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool
        Whether to reload the dataset. Default: False
    verbose : bool
        Whether to print out progress information. Default: True.
    transform : callable, optional
        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
        a transformed version. The :class:`~dgl.DGLGraph` object will be
        transformed before every access.

    Attributes
    ----------
    num_classes : int
        Number of classes to predict
    predict_category : str
        All the labels of the entities in ``predict_category``

    Examples
    --------
    >>> dataset = dgl.data.rdf.BGSDataset()
    >>> graph = dataset[0]
    >>> category = dataset.predict_category
    >>> num_classes = dataset.num_classes
    >>>
    >>> train_mask = g.nodes[category].data['train_mask']
    >>> test_mask = g.nodes[category].data['test_mask']
    >>> label = g.nodes[category].data['label']
    """

    entity_prefix = "http://data.bgs.ac.uk/"
    status_prefix = "http://data.bgs.ac.uk/ref/CurrentStatus"
    relation_prefix = "http://data.bgs.ac.uk/ref"

    def __init__(
        self,
        print_every=10000,
        insert_reverse=True,
        raw_dir=None,
        force_reload=False,
        verbose=True,
        transform=None,
    ):
        import rdflib as rdf

        url = _get_dgl_url("dataset/rdf/bgs-hetero.zip")
        name = "bgs-hetero"
        predict_category = "Lexicon/NamedRockUnit"
        self.lith = rdf.term.URIRef(
            "http://data.bgs.ac.uk/ref/Lexicon/hasLithogenesis"
        )
        super(BGSDataset, self).__init__(
            name,
            url,
            predict_category,
            print_every=print_every,
            insert_reverse=insert_reverse,
            raw_dir=raw_dir,
            force_reload=force_reload,
            verbose=verbose,
            transform=transform,
        )



    def __getitem__(self, idx):
        r"""Gets the graph object

        Parameters
        -----------
        idx: int
            Item index, BGSDataset has only one graph object

        Return
        -------
        :class:`dgl.DGLGraph`

            The graph contains:

            - ``ndata['train_mask']``: mask for training node set
            - ``ndata['test_mask']``: mask for testing node set
            - ``ndata['label']``: node labels
        """
        return super(BGSDataset, self).__getitem__(idx)




    def __len__(self):
        r"""The number of graphs in the dataset.

        Return
        -------
        int
        """
        return super(BGSDataset, self).__len__()


    def parse_entity(self, term):
        import rdflib as rdf

        if isinstance(term, rdf.Literal):
            return None
        elif isinstance(term, rdf.BNode):
            return None
        entstr = str(term)
        if entstr.startswith(self.status_prefix):
            return None
        if entstr.startswith(self.entity_prefix):
            sp = entstr.split("/")
            if len(sp) != 7:
                return None
            # instance
            cls = "%s/%s" % (sp[4], sp[5])
            inst = sp[6]
            return Entity(e_id=inst, cls=cls)
        else:
            return None

    def parse_relation(self, term):
        if term == self.lith:
            return None
        relstr = str(term)
        if relstr.startswith(self.relation_prefix):
            sp = relstr.split("/")
            if len(sp) < 6:
                return None
            assert len(sp) == 6, relstr
            cls = "%s/%s" % (sp[4], sp[5])
            return Relation(cls=cls)
        else:
            relstr = relstr.replace(".", "_")
            return Relation(cls=relstr)

    def process_tuple(self, raw_tuple, sbj, rel, obj):
        if sbj is None or rel is None or obj is None:
            return None
        return (sbj, rel, obj)

    def process_idx_file_line(self, line):
        _, rock, label = line.strip().split("\t")
        return rock, label





class AMDataset(RDFGraphDataset):
    """AM dataset. for node classification task

    Namespace convention:

    - Instance: ``http://purl.org/collections/nl/am/<type>-<id>``
    - Relation: ``http://purl.org/collections/nl/am/<name>``

    We ignored all literal nodes and the relations connecting them in the
    output graph.

    AM dataset statistics:

    - Nodes: 881680
    - Edges: 5668682 (including reverse edges)
    - Target Category: proxy
    - Number of Classes: 11
    - Label Split:

        - Train: 802
        - Test: 198

    Parameters
    -----------
    print_every : int
        Preprocessing log for every X tuples. Default: 10000.
    insert_reverse : bool
        If true, add reverse edge and reverse relations to the final graph. Default: True.
    raw_dir : str
        Raw file directory to download/contains the input data directory.
        Default: ~/.dgl/
    force_reload : bool
        Whether to reload the dataset. Default: False
    verbose : bool
        Whether to print out progress information. Default: True.
    transform : callable, optional
        A transform that takes in a :class:`~dgl.DGLGraph` object and returns
        a transformed version. The :class:`~dgl.DGLGraph` object will be
        transformed before every access.

    Attributes
    ----------
    num_classes : int
        Number of classes to predict
    predict_category : str
        The entity category (node type) that has labels for prediction

    Examples
    --------
    >>> dataset = dgl.data.rdf.AMDataset()
    >>> graph = dataset[0]
    >>> category = dataset.predict_category
    >>> num_classes = dataset.num_classes
    >>>
    >>> train_mask = g.nodes[category].data['train_mask']
    >>> test_mask = g.nodes[category].data['test_mask']
    >>> label = g.nodes[category].data['label']
    """

    entity_prefix = "http://purl.org/collections/nl/am/"
    relation_prefix = entity_prefix

    def __init__(
        self,
        print_every=10000,
        insert_reverse=True,
        raw_dir=None,
        force_reload=False,
        verbose=True,
        transform=None,
    ):
        import rdflib as rdf

        self.objectCategory = rdf.term.URIRef(
            "http://purl.org/collections/nl/am/objectCategory"
        )
        self.material = rdf.term.URIRef(
            "http://purl.org/collections/nl/am/material"
        )
        url = _get_dgl_url("dataset/rdf/am-hetero.zip")
        name = "am-hetero"
        predict_category = "proxy"
        super(AMDataset, self).__init__(
            name,
            url,
            predict_category,
            print_every=print_every,
            insert_reverse=insert_reverse,
            raw_dir=raw_dir,
            force_reload=force_reload,
            verbose=verbose,
            transform=transform,
        )



    def __getitem__(self, idx):
        r"""Gets the graph object

        Parameters
        -----------
        idx: int
            Item index, AMDataset has only one graph object

        Return
        -------
        :class:`dgl.DGLGraph`

            The graph contains:

            - ``ndata['train_mask']``: mask for training node set
            - ``ndata['test_mask']``: mask for testing node set
            - ``ndata['label']``: node labels
        """
        return super(AMDataset, self).__getitem__(idx)




    def __len__(self):
        r"""The number of graphs in the dataset.

        Return
        -------
        int
        """
        return super(AMDataset, self).__len__()


    def parse_entity(self, term):
        import rdflib as rdf

        if isinstance(term, rdf.Literal):
            return None
        elif isinstance(term, rdf.BNode):
            return Entity(e_id=str(term), cls="_BNode")
        entstr = str(term)
        if entstr.startswith(self.entity_prefix):
            sp = entstr.split("/")
            assert len(sp) == 7, entstr
            spp = sp[6].split("-")
            if len(spp) == 2:
                # instance
                cls, inst = spp
            else:
                cls = "TYPE"
                inst = spp
            return Entity(e_id=inst, cls=cls)
        else:
            return None

    def parse_relation(self, term):
        if term == self.objectCategory or term == self.material:
            return None
        relstr = str(term)
        if relstr.startswith(self.relation_prefix):
            sp = relstr.split("/")
            assert len(sp) == 7, relstr
            cls = sp[6]
            return Relation(cls=cls)
        else:
            relstr = relstr.replace(".", "_")
            return Relation(cls=relstr)

    def process_tuple(self, raw_tuple, sbj, rel, obj):
        if sbj is None or rel is None or obj is None:
            return None
        return (sbj, rel, obj)

    def process_idx_file_line(self, line):
        proxy, _, label = line.strip().split("\t")
        return proxy, label