def _convert_graph_from_ogb()

in hugegraph-ml/src/hugegraph_ml/data/hugegraph2dgl.py [0:0]


    def _convert_graph_from_ogb(vertices, edges, feat_key, year_key, weight_key):
        if len(vertices) == 0:
            warnings.warn("This graph has no vertices", Warning)
            return dgl.graph(())
        vertex_ids = [v["id"] for v in vertices]
        vertex_id_to_idx = {vertex_id: idx for idx, vertex_id in enumerate(vertex_ids)}
        src_idx = [vertex_id_to_idx[e["outV"]] for e in edges]
        dst_idx = [vertex_id_to_idx[e["inV"]] for e in edges]
        graph_dgl = dgl.graph((src_idx, dst_idx))
        if feat_key and feat_key in vertices[0]["properties"]:
            node_feats = [
                v["properties"][feat_key]
                for v in vertices[0 : graph_dgl.number_of_nodes()]
            ]
            graph_dgl.ndata["feat"] = torch.tensor(node_feats, dtype=torch.float32)
        if year_key and year_key in edges[0]["properties"]:
            year = [e["properties"][year_key] for e in edges]
            graph_dgl.edata["year"] = torch.tensor(year, dtype=torch.int64)
        if weight_key and weight_key in edges[0]["properties"]:
            weight = [e["properties"][weight_key] for e in edges]
            graph_dgl.edata["weight"] = torch.tensor(weight, dtype=torch.int64)

        return graph_dgl, vertex_id_to_idx