<commit_before>from __future__ import print_function

import argparse
import random
import time

import networkx as nx
import numpy as np
from gensim.models import Word2Vec


class Node2Vec(object):
    def __init__(self, is_directed, preprocessing, is_weighted, p, q, walk_length, num_walks, dimensions, window_size,
                 workers, iterations):

        # graph properties
        self.is_directed = is_directed
        self.preprocessing = preprocessing
        self.is_weighted = is_weighted
        self.G = None
        self.alias_nodes = None
        self.alias_edges = None

        # walk properties
        self.p = p
        self.q = q
        self.walk_length = walk_length
        self.num_walks = num_walks

        # learning properties
        self.dimensions = dimensions
        self.window_size = window_size
        self.workers = workers
        self.iter = iterations

    def read_graph(self, nx_g):

        if self.is_weighted:

            self.G = nx.read_edgelist(nx_g, data=(('weight', float),), create_using=nx.DiGraph(), edgetype=str)

        else:

            self.G = nx.read_edgelist(nx_g, create_using=nx.DiGraph(), edgetype=str)

        for edge in self.G.edges():
            self.G[edge[0]][edge[1]]['weight'] = 1

        if not self.is_directed:
            self.G = self.G.to_undirected()

    def node2vec_walk(self, start_node):
        """
        Simulate a random walk starting from start node.
        """

        G = self.G

        walk = [start_node]

        while len(walk) < self.walk_length:

            cur = walk[-1]
            cur_nbrs = sorted(G.neighbors(cur))

            if len(cur_nbrs) > 0:

                if self.preprocessing:

                    alias_nodes = self.alias_nodes
                    alias_edges = self.alias_edges

                    if len(walk) == 1:  # first step of the walk, no previous node

                        walk.append(cur_nbrs[Node2Vec.alias_draw(alias_nodes[cur][0], alias_nodes[cur][1])])

                    else:
                        prev = walk[-2]

                        next_node = cur_nbrs[Node2Vec.alias_draw(alias_edges[(prev, cur)][0],
                                                                 alias_edges[(prev, cur)][1])]

                        walk.append(next_node)

                else:

                    p = self.p
                    q = self.q
                    G = self.G

                    unnormalized_probs = []

                    if len(walk) == 1:  # first step of the walk, no previous node

                        for dst_nbr in cur_nbrs:
                            unnormalized_probs.append(G[cur][dst_nbr]['weight'])

                        norm_const = sum(unnormalized_probs)

                        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

                        next_node = cur_nbrs[np.random.multinomial(1, normalized_probs).argmax()]

                        walk.append(next_node)

                    else:

                        prev = walk[-2]

                        for dst_nbr in cur_nbrs:

                            if dst_nbr == prev:

                                unnormalized_probs.append(G[cur][dst_nbr]['weight'] / p)

                            elif G.has_edge(dst_nbr, prev):
                                unnormalized_probs.append(G[cur][dst_nbr]['weight'])
                            else:
                                unnormalized_probs.append(G[cur][dst_nbr]['weight'] / q)

                        norm_const = sum(unnormalized_probs)
                        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

                        next_node = cur_nbrs[np.random.multinomial(1, normalized_probs).argmax()]

                        walk.append(next_node)
            else:
                break

        return walk

    def learn_embeddings(self, output, output_format='binary'):
        """
        Learn embeddings by optimizing the Skipgram objective using SGD.
        """

        self._simulate_walks()  # simulate random walks

        model = Word2Vec(self._walks, size=self.dimensions, window=self.window_size, min_count=0,
                         workers=self.workers, iter=self.iter, negative=25, sg=1)

        print("defined model using w2v")

        is_binary = output_format != 'text'
        model.wv.save_word2vec_format(output, binary=is_binary)

        actual_format = 'text' if output_format == 'text' else 'binary'
        print("saved model in word2vec %s format" % actual_format)

        return

    def get_alias_edge(self, src, dst):
        """
        Get the alias edge setup lists for a given edge.
        """
        G = self.G
        p = self.p
        q = self.q

        unnormalized_probs = []
        for dst_nbr in sorted(G.neighbors(dst)):
            if dst_nbr == src:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / p)
            elif G.has_edge(dst_nbr, src):
                unnormalized_probs.append(G[dst][dst_nbr]['weight'])
            else:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / q)
        norm_const = sum(unnormalized_probs)
        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

        return Node2Vec.alias_setup(normalized_probs)

    def _simulate_walks(self):

        """
        Simulate random walks from each node.
        """
        G = self.G
        nodes = G.nodes()

        self._walks = []

        print('Walk iteration:')

        for walk_iter in range(self.num_walks):

            print(str(walk_iter + 1), '/', str(self.num_walks))
            random.shuffle(nodes)
<commit_msg>list cast to G.nodes()<commit_after>from __future__ import print_function

import argparse
import random
import time

import networkx as nx
import numpy as np
from gensim.models import Word2Vec


class Node2Vec(object):
    def __init__(self, is_directed, preprocessing, is_weighted, p, q, walk_length, num_walks, dimensions, window_size,
                 workers, iterations):

        # graph properties
        self.is_directed = is_directed
        self.preprocessing = preprocessing
        self.is_weighted = is_weighted
        self.G = None
        self.alias_nodes = None
        self.alias_edges = None

        # walk properties
        self.p = p
        self.q = q
        self.walk_length = walk_length
        self.num_walks = num_walks

        # learning properties
        self.dimensions = dimensions
        self.window_size = window_size
        self.workers = workers
        self.iter = iterations

    def read_graph(self, nx_g):

        if self.is_weighted:

            self.G = nx.read_edgelist(nx_g, data=(('weight', float),), create_using=nx.DiGraph(), edgetype=str)

        else:

            self.G = nx.read_edgelist(nx_g, create_using=nx.DiGraph(), edgetype=str)

        for edge in self.G.edges():
            self.G[edge[0]][edge[1]]['weight'] = 1

        if not self.is_directed:
            self.G = self.G.to_undirected()

    def node2vec_walk(self, start_node):
        """
        Simulate a random walk starting from start node.
        """

        G = self.G

        walk = [start_node]

        while len(walk) < self.walk_length:

            cur = walk[-1]
            cur_nbrs = sorted(G.neighbors(cur))

            if len(cur_nbrs) > 0:

                if self.preprocessing:

                    alias_nodes = self.alias_nodes
                    alias_edges = self.alias_edges

                    if len(walk) == 1:  # first step of the walk, no previous node

                        walk.append(cur_nbrs[Node2Vec.alias_draw(alias_nodes[cur][0], alias_nodes[cur][1])])

                    else:
                        prev = walk[-2]

                        next_node = cur_nbrs[Node2Vec.alias_draw(alias_edges[(prev, cur)][0],
                                                                 alias_edges[(prev, cur)][1])]

                        walk.append(next_node)

                else:

                    p = self.p
                    q = self.q
                    G = self.G

                    unnormalized_probs = []

                    if len(walk) == 1:  # first step of the walk, no previous node

                        for dst_nbr in cur_nbrs:
                            unnormalized_probs.append(G[cur][dst_nbr]['weight'])

                        norm_const = sum(unnormalized_probs)

                        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

                        next_node = cur_nbrs[np.random.multinomial(1, normalized_probs).argmax()]

                        walk.append(next_node)

                    else:

                        prev = walk[-2]

                        for dst_nbr in cur_nbrs:

                            if dst_nbr == prev:

                                unnormalized_probs.append(G[cur][dst_nbr]['weight'] / p)

                            elif G.has_edge(dst_nbr, prev):
                                unnormalized_probs.append(G[cur][dst_nbr]['weight'])
                            else:
                                unnormalized_probs.append(G[cur][dst_nbr]['weight'] / q)

                        norm_const = sum(unnormalized_probs)
                        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

                        next_node = cur_nbrs[np.random.multinomial(1, normalized_probs).argmax()]

                        walk.append(next_node)
            else:
                break

        return walk

    def learn_embeddings(self, output, output_format='binary'):
        """
        Learn embeddings by optimizing the Skipgram objective using SGD.
        """

        self._simulate_walks()  # simulate random walks

        model = Word2Vec(self._walks, size=self.dimensions, window=self.window_size, min_count=0,
                         workers=self.workers, iter=self.iter, negative=25, sg=1)

        print("defined model using w2v")

        is_binary = output_format != 'text'
        model.wv.save_word2vec_format(output, binary=is_binary)

        actual_format = 'text' if output_format == 'text' else 'binary'
        print("saved model in word2vec %s format" % actual_format)

        return

    def get_alias_edge(self, src, dst):
        """
        Get the alias edge setup lists for a given edge.
        """
        G = self.G
        p = self.p
        q = self.q

        unnormalized_probs = []
        for dst_nbr in sorted(G.neighbors(dst)):
            if dst_nbr == src:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / p)
            elif G.has_edge(dst_nbr, src):
                unnormalized_probs.append(G[dst][dst_nbr]['weight'])
            else:
                unnormalized_probs.append(G[dst][dst_nbr]['weight'] / q)
        norm_const = sum(unnormalized_probs)
        normalized_probs = [float(u_prob) / norm_const for u_prob in unnormalized_probs]

        return Node2Vec.alias_setup(normalized_probs)

    def _simulate_walks(self):

        """
        Simulate random walks from each node.
        """
        G = self.G
        nodes = list(G.nodes())

        self._walks = []

        print('Walk iteration:')

        for walk_iter in range(self.num_walks):

            print(str(walk_iter + 1), '/', str(self.num_walks))
            random.shuffle(nodes)

{'fraction_non_alphanumeric': 0.09424767209120181, 'fraction_numerical': 0.005221477678182926, 'mean_word_length': 1.6064867316851894, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3587828', 'n_tokens_mistral': 3495, 'n_tokens_neox': 3248, 'n_words': 693}