<commit_before>#!/usr/bin/env python

import argparse
import os
import pickle
from gensim.models.word2vec import Word2Vec
import csv
import sys
import numpy as np
import tensorflow as tf
from projlearn import *

parser = argparse.ArgumentParser(description='Prediction.')
parser.add_argument('--w2v',    required=True, type=argparse.FileType('rb'))
parser.add_argument('--kmeans', default='kmeans.pickle', nargs='?', help='Path to k-means.pickle.')
parser.add_argument('--model',  default='baseline', nargs='?', choices=MODELS.keys(), help='The model.')
parser.add_argument('--path',   default='', nargs='?', help='The path to the model dump.')
parser.add_argument('output',   type=argparse.FileType('wb'), help='The output npz file.')
args = parser.parse_args()

w2v = Word2Vec.load_word2vec_format(args.w2v, binary=True, unicode_errors='ignore')
w2v.init_sims(replace=True)

print('Using %d word2vec dimensions from "%s".' % (w2v.layer1_size, args.w2v.name), flush=True, file=sys.stderr)

kmeans = pickle.load(open(args.kmeans, 'rb'))

print('The number of clusters is %d.' % kmeans.n_clusters, flush=True, file=sys.stderr)

reader = csv.reader(sys.stdin, delimiter='\t', quoting=csv.QUOTE_NONE)

X_all, Y_all = [], []

for row in reader:
    X_all.append(w2v.wv.vocab[row[0]].index)
    Y_all.append(w2v.wv.vocab[row[1]].index)

X_all, Y_all = w2v.wv.syn0[X_all], w2v.wv.syn0[Y_all]

offsets = Y_all - X_all

X_clusters_list = list(enumerate(kmeans.predict(offsets)))

X_clusters = {}

for cluster in range(kmeans.n_clusters):
    X_clusters[cluster] = [i for i, c in X_clusters_list if c == cluster]

model = MODELS[args.model](x_size=w2v.layer1_size, y_size=w2v.layer1_size, w_stddev=0, lambda_=0)

Y_hat_all = np.empty(X_all.shape)

for cluster, indices in X_clusters.items():
    with tf.Session() as sess:
        saver = tf.train.Saver()

        saver.restore(sess, os.path.join(args.path, '%s.k%d.trained') % (args.model, cluster + 1))

        Y_hat = sess.run(model.Y_hat, feed_dict={model.X: X_all[indices]})

        for i, j in enumerate(indices):
            Y_hat_all[j] = Y_hat[i]

np.savez_compressed(args.output, Y_hat_all=Y_hat_all)<commit_msg>Use grouping and HDF5<commit_after>#!/usr/bin/env python

import argparse
import os
import pickle
from gensim.models.word2vec import Word2Vec
import sys
import csv
from itertools import zip_longest
import numpy as np
import tensorflow as tf
from projlearn import *
import h5py

os.environ['CUDA_VISIBLE_DEVICES'] = ''

def grouper(n, iterable, fillvalue=None):
    args = [iter(iterable)] * n
    return zip_longest(fillvalue=fillvalue, *args)

parser = argparse.ArgumentParser(description='Prediction.')
parser.add_argument('--w2v',    required=True, type=argparse.FileType('rb'))
parser.add_argument('--kmeans', default='kmeans.pickle', nargs='?', help='Path to k-means.pickle.')
parser.add_argument('--model',  default='baseline', nargs='?', choices=MODELS.keys(), help='The model.')
parser.add_argument('--path',   default='', nargs='?', help='The path to the model dump.')
parser.add_argument('--slices', default=100000, type=int, help='The slice size.')
parser.add_argument('hdf5')
args = parser.parse_args()

w2v = Word2Vec.load_word2vec_format(args.w2v, binary=True, unicode_errors='ignore')
w2v.init_sims(replace=True)

print('Using %d word2vec dimensions from "%s".' % (w2v.layer1_size, args.w2v.name), flush=True, file=sys.stderr)

kmeans = pickle.load(open(args.kmeans, 'rb'))

print('The number of clusters is %d.' % kmeans.n_clusters, flush=True, file=sys.stderr)

model = MODELS[args.model](x_size=w2v.layer1_size, y_size=w2v.layer1_size, w_stddev=0, lambda_=0)

reader = csv.reader(sys.stdin, delimiter='\t', quoting=csv.QUOTE_NONE)

try:
    os.remove(args.hdf5)
except OSError:
    pass

with h5py.File(args.hdf5, 'w') as f:
    dset = f.create_dataset('hyperstar', (0, w2v.layer1_size), maxshape=(None, w2v.layer1_size), chunks=True)

    for rows in grouper(args.slices, reader):
        X_all, Y_all = [], []

        for row in rows:
            if row is None:
                continue

            X_all.append(w2v.wv.vocab[row[0]].index)
            Y_all.append(w2v.wv.vocab[row[1]].index)

        X_all, Y_all = w2v.wv.syn0[X_all], w2v.wv.syn0[Y_all]

        offsets = Y_all - X_all

        X_clusters_list = list(enumerate(kmeans.predict(offsets)))

        X_clusters = {}

        for cluster in range(kmeans.n_clusters):
            X_clusters[cluster] = [i for i, c in X_clusters_list if c == cluster]

        Y_hat_all = np.empty(X_all.shape)

        for cluster, indices in X_clusters.items():
            with tf.Session() as sess:
                saver = tf.train.Saver()

                saver.restore(sess, os.path.join(args.path, '%s.k%d.trained') % (args.model, cluster + 1))

                Y_hat = sess.run(model.Y_hat, feed_dict={model.X: X_all[indices]})

                for i, j in enumerate(indices):
                    Y_hat_all[j] = Y_hat[i]

        dset.resize(dset.shape[0] + Y_hat_all.shape[0], axis=0)

        dset[-Y_hat_all.shape[0]:] = Y_hat_all

        print('The number of rows is %d.' % dset.shape[0], flush=True, file=sys.stderr)
{'fraction_non_alphanumeric': 0.1531461761858664, 'fraction_numerical': 0.014714424007744434, 'mean_word_length': 4.020408163265306, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2909837', 'n_tokens_mistral': 2013, 'n_tokens_neox': 1941, 'n_words': 381}