#!/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 *

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('output',   type=argparse.FileType('wb'), help='Output 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)

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)

for s, rows in enumerate(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]

    np.save(args.output, Y_hat_all, allow_pickle=False)

    print('%d slices done.' % (s + 1), flush=True, file=sys.stderr)

@@ -222,16 +222,42 @@
 import *
+%0Afrom gzip import GzipFile
 %0A%0Aos.env
@@ -928,24 +928,90 @@
 ice size.')%0A
+parser.add_argument('--gzip', default=False, action='store_true')%0A
 parser.add_a
@@ -1643,16 +1643,106 @@
 _NONE)%0A%0A
+with args.output if not args.gzip else GzipFile(fileobj=args.output, mode='wb') as f:%0A    
 for s, r
@@ -1789,16 +1789,20 @@
 ader)):%0A
+    
     X_al
@@ -1820,24 +1820,28 @@
 %5B%5D, %5B%5D%0A%0A    
+ 
+   
 for row in r
@@ -1853,16 +1853,20 @@
         
+   
+ 
 if row i
@@ -1885,16 +1885,20 @@
         
+    
 continue
@@ -1903,24 +1903,28 @@
 ue%0A%0A        
+ 
+   
 X_all.append
@@ -1960,16 +1960,20 @@
         
+   
+ 
 Y_all.ap
@@ -2006,16 +2006,20 @@
 index)%0A%0A
+    
     X_al
@@ -2073,16 +2073,20 @@
 l%5D%0A%0A    
+ 
+   
 offsets 
@@ -2102,16 +2102,20 @@
  X_all%0A%0A
+    
     X_cl
@@ -2166,24 +2166,28 @@
 offsets)))%0A%0A
+    
     X_cluste
@@ -2195,24 +2195,28 @@
 s = %7B%7D%0A%0A    
+ 
+   
 for cluster 
@@ -2244,32 +2244,36 @@
 sters):%0A        
+   
+ 
 X_clusters%5Bclust
@@ -2327,16 +2327,20 @@
 uster%5D%0A%0A
+    
     Y_ha
@@ -2370,24 +2370,28 @@
 shape)%0A%0A    
+ 
+   
 for cluster,
@@ -2430,16 +2430,20 @@
         
+   
+ 
 with tf.
@@ -2469,24 +2469,28 @@
             
+    
 saver = tf.t
@@ -2511,24 +2511,28 @@
             
+ 
+   
 saver.restor
@@ -2619,24 +2619,28 @@
             
+   
+ 
 Y_hat = sess
@@ -2703,24 +2703,28 @@
             
+    
 for i, j in 
@@ -2751,32 +2751,36 @@
                 
+ 
+   
 Y_hat_all%5Bj%5D = Y
@@ -2788,16 +2788,20 @@
 hat%5Bi%5D%0A%0A
+    
     np.s
@@ -2804,27 +2804,17 @@
 np.save(
-args.output
+f
 , Y_hat_
@@ -2835,24 +2835,28 @@
 kle=False)%0A%0A
+    
     print('%25

Use Gzip