Function main

examples/node2vec_wiki.py:50–76

(smoke=False, show=True)

def main(smoke=False, show=True):
    graph_path = SMOKE_GRAPH_PATH if smoke else WIKI_GRAPH_PATH
    graph = nx.read_edgelist(
        str(graph_path),
        create_using=nx.DiGraph(),
        nodetype=None,
        data=[("weight", int)],
    )

    model = Node2Vec(
        graph,
        walk_length=3 if smoke else 10,
        num_walks=2 if smoke else 80,
        p=0.25,
        q=4,
        workers=1,
        use_rejection_sampling=False,
    )
    model.train(window_size=2 if smoke else 5, iter=1 if smoke else 3, workers=1)
    embeddings = model.get_embeddings()
    assert len(embeddings) > 0

    if not smoke:
        evaluate_embeddings(embeddings, WIKI_LABEL_PATH)
        plot_embeddings(embeddings, WIKI_LABEL_PATH, show=show)

    return embeddings


if __name__ == "__main__":