Function main

examples/sdne_wiki.py:56–78

(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 = SDNE(graph, hidden_size=[8, 4] if smoke else [256, 128])
    model.train(
        batch_size=2 if smoke else 3000,
        epochs=1 if smoke else 40,
        verbose=0 if smoke else 2,
    )
    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__":