def convert_data()

in hugegraph-ml/src/hugegraph_ml/models/bgnn.py [0:0]

• 26 lines of code
• 7 McCabe index (conditional complexity)

def convert_data(g):
    retrieved_tensor = g.ndata["feat"]
    retrieved_np = retrieved_tensor.numpy()
    retrieved_str = retrieved_np.astype(str)
    X = pd.DataFrame(retrieved_str)

    retrieved_y_tensor = g.ndata["class"]
    retrieved_y_np = retrieved_y_tensor.numpy()
    y = pd.DataFrame(retrieved_y_np)

    retrieved_cat_features_tensor = g.ndata["cat_features"][0]
    cat_features = retrieved_cat_features_tensor.numpy()

    train_mask = g.ndata["train_mask"].numpy().tolist()
    val_mask = g.ndata["val_mask"].numpy().tolist()
    test_mask = g.ndata["test_mask"].numpy().tolist()
    masks = {
        "0": {
            "train": [i for i, v in enumerate(train_mask) if v == 1],
            "val": [i for i, v in enumerate(val_mask) if v == 1],
            "test": [i for i, v in enumerate(test_mask) if v == 1],
        }
    }

    # graph, X, y, cat_features, masks = read_input(input_folder)
    train_mask, val_mask, test_mask = (
        masks["0"]["train"],
        masks["0"]["val"],
        masks["0"]["test"],
    )

    return X, y, cat_features, train_mask, val_mask, test_mask