def _convert_split_edge_from_ogb()

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

• 65 lines of code
• 13 McCabe index (conditional complexity)

    def _convert_split_edge_from_ogb(edges, vertex_id_to_idx):
        train_edge_list = []
        train_year_list = []
        train_weight_list = []
        valid_edge_list = []
        valid_year_list = []
        valid_weight_list = []
        valid_edge_neg_list = []
        test_edge_list = []
        test_year_list = []
        test_weight_list = []
        test_edge_neg_list = []

        for edge in edges:
            if edge["properties"]["train_edge_mask"] == 1:
                train_edge_list.append(
                    [vertex_id_to_idx[edge["outV"]], vertex_id_to_idx[edge["inV"]]]
                )
            if edge["properties"]["train_year_mask"] != -1:
                train_year_list.append(edge["properties"]["train_year_mask"])
            if edge["properties"]["train_weight_mask"] != -1:
                train_weight_list.append(edge["properties"]["train_weight_mask"])

            if edge["properties"]["valid_edge_mask"] == 1:
                valid_edge_list.append(
                    [vertex_id_to_idx[edge["outV"]], vertex_id_to_idx[edge["inV"]]]
                )
            if edge["properties"]["valid_year_mask"] != -1:
                valid_year_list.append(edge["properties"]["valid_year_mask"])
            if edge["properties"]["valid_weight_mask"] != -1:
                valid_weight_list.append(edge["properties"]["valid_weight_mask"])
            if edge["properties"]["valid_edge_neg_mask"] == 1:
                valid_edge_neg_list.append(
                    [vertex_id_to_idx[edge["outV"]], vertex_id_to_idx[edge["inV"]]]
                )

            if edge["properties"]["test_edge_mask"] == 1:
                test_edge_list.append(
                    [vertex_id_to_idx[edge["outV"]], vertex_id_to_idx[edge["inV"]]]
                )
            if edge["properties"]["test_year_mask"] != -1:
                test_year_list.append(edge["properties"]["test_year_mask"])
            if edge["properties"]["test_weight_mask"] != -1:
                test_weight_list.append(edge["properties"]["test_weight_mask"])
            if edge["properties"]["test_edge_neg_mask"] == 1:
                test_edge_neg_list.append(
                    [vertex_id_to_idx[edge["outV"]], vertex_id_to_idx[edge["inV"]]]
                )

        split_edge = {
            "train": {
                "edge": torch.tensor(train_edge_list),
                "weight": torch.tensor(train_weight_list),
                "year": torch.tensor(train_year_list),
            },
            "valid": {
                "edge": torch.tensor(valid_edge_list),
                "weight": torch.tensor(valid_weight_list),
                "year": torch.tensor(valid_year_list),
                "edge_neg": torch.tensor(valid_edge_neg_list),
            },
            "test": {
                "edge": torch.tensor(test_edge_list),
                "weight": torch.tensor(test_weight_list),
                "year": torch.tensor(test_year_list),
                "edge_neg": torch.tensor(test_edge_neg_list),
            },
        }

        return split_edge