myutils.py

import sys
import javalang
from timeit import default_timer as timer
import keras
import numpy as np
import tensorflow as tf
import networkx as nx
import random
import pickle
import traceback

# do NOT import keras in this header area, it will break predict.py
# instead, import keras as needed in each function

# TODO refactor this so it imports in the necessary functions
dataprep = '/nfs/projects/firstwords/data/standard'
sys.path.append(dataprep)
import tokenizer

start = 0
end = 0

def init_tf(gpu, horovod=False):
    from keras.backend.tensorflow_backend import set_session
    
    config = tf.ConfigProto()
    config = tf.ConfigProto(log_device_placement=False)
    config.gpu_options.allow_growth = True
    config.gpu_options.visible_device_list = gpu

    set_session(tf.Session(config=config))

def prep(msg):
    global start
    statusout(msg)
    start = timer()

def statusout(msg):
    sys.stdout.write(msg)
    sys.stdout.flush()

def drop():
    global start
    global end
    end = timer()
    sys.stdout.write('done, %s seconds.\n' % (round(end - start, 2)))
    sys.stdout.flush()

def index2word(tok):
	i2w = {}
	for word, index in tok.w2i.items():
		i2w[index] = word

	return i2w

def seq2sent(seq, tokenizer):
    sent = []
    check = index2word(tokenizer)
    for i in seq:
        sent.append(check[i])

    return(' '.join(sent))
            
class batch_gen(keras.utils.Sequence):
    def __init__(self, seqdata, firstwords, tt, config, training=True):
        self.comvocabsize = config['comvocabsize']
        self.tt = tt
        self.firstwords = firstwords['%sfw' % (tt)]
        self.batch_size = config['batch_size']
        self.seqdata = seqdata
        self.allfids = list(seqdata['dt%s' % (tt)].keys())
        self.num_inputs = config['num_input']
        self.config = config
        self.training = training
        
        random.shuffle(self.allfids) # actually, might need to sort allfids to ensure same order

    def __getitem__(self, idx):
        start = (idx*self.batch_size)
        end = self.batch_size*(idx+1)
        batchfids = self.allfids[start:end]
        return self.make_batch(batchfids)

    def make_batch(self, batchfids):
        if self.config['batch_maker'] == 'datsonly':
            return self.divideseqs_datsonly(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'seqsonly':
            return self.divideseqs(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'ast':
            return self.divideseqs_ast(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'ast_threed':
            return self.divideseqs_ast_threed(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'threed':
            return self.divideseqs_threed(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'graphast':
            return self.divideseqs_graphast(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'graphast_threed':
            return self.divideseqs_graphast_threed(batchfids, self.seqdata, self.comvocabsize, self.tt)
        elif self.config['batch_maker'] == 'pathast_threed':
            return self.divideseqs_pathast_threed(batchfids, self.seqdata, self.comvocabsize, self.tt)
        else:
            return None

    def __len__(self):
        #if self.num_inputs == 4:
        return int(np.ceil(len(list(self.seqdata['dt%s' % (self.tt)]))/self.batch_size))
        #else:
        #    return int(np.ceil(len(list(self.seqdata['d%s' % (self.tt)]))/self.batch_size))

    def on_epoch_end(self):
        random.shuffle(self.allfids)

    def divideseqs(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils
        
        smlseqs = list()
        comouts = list()
        
        fiddat = dict()

        for fid in batchfids:

            wsmlseq = seqdata['s%s' % (tt)][fid]

            if not self.training:
                fiddat[fid] = [wsmlseq]
            else:
                smlseqs.append(wsmlseq)
                comout = keras.utils.to_categorical(self.firstwords[fid],
                                                    num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        smlseqs = np.asarray(smlseqs)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            return [[smlseqs], comouts]

    def divideseqs_datsonly(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        datseqs = list()
        comouts = list()

        fiddat = dict()
        
        for fid in batchfids:
            wdatseq = seqdata['dt%s' % (tt)][fid]
            wdatseq = wdatseq[:self.config['tdatlen']]
            
            if not self.training:
                fiddat[fid] = [wdatseq]
            else:
                datseqs.append(wdatseq)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))
        datseqs = np.asarray(datseqs)
        comouts = np.asarray(comouts)
        if not self.training:
            return fiddat
        else:
            return [[datseqs], comouts]
        return

    def divideseqs_ast(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils
        
        datseqs = list()
        smlseqs = list()
        comouts = list()
        
        fiddat = dict()

        for fid in batchfids:

            wdatseq = seqdata['dt%s' % (tt)][fid]
            wsmlseq = seqdata['s%s' % (tt)][fid]

            wdatseq = wdatseq[:self.config['tdatlen']]

            if not self.training:
                fiddat[fid] = [wdatseq, wsmlseq]
            else:
                datseqs.append(wdatseq)
                smlseqs.append(wsmlseq)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        datseqs = np.asarray(datseqs)
        smlseqs = np.asarray(smlseqs)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            return [[datseqs, smlseqs], comouts]

    def divideseqs_ast_threed(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        tdatseqs = list()
        sdatseqs = list()
        smlseqs = list()
        comouts = list()

        fiddat = dict()

        for fid in batchfids:
            wtdatseq = seqdata['dt%s' % (tt)][fid]
            wsdatseq = seqdata['ds%s' % (tt)][fid]
            wsmlseq = seqdata['s%s' % (tt)][fid]
            
            wtdatseq = wtdatseq[:self.config['tdatlen']]

            newlen = self.config['sdatlen'] - len(wsdatseq)
            if newlen < 0:
                newlen=0
            wsdatseq = wsdatseq.tolist()
            for k in range(newlen):
                wsdatseq.append(np.zeros(self.config['stdatlen']))
            for i in range(0, len(wsdatseq)):
                wsdatseq[i] = np.array(wsdatseq[i])[:self.config['stdatlen']]
            wsdatseq = np.asarray(wsdatseq)
            wsdatseq = wsdatseq[:self.config['sdatlen'], :]
            
            wsmlseq = wsmlseq[:self.config['smllen']]
            
            if not self.training:
                fiddat[fid] = [wtdatseq, wsdatseq, wsmlseq]
            else:
                tdatseqs.append(wtdatseq)
                sdatseqs.append(wsdatseq)
                smlseqs.append(wsmlseq)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        tdatseqs = np.asarray(tdatseqs)
        sdatseqs = np.asarray(sdatseqs)
        smlseqs = np.asarray(smlseqs)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            if self.config['num_output'] == 2:
                return [[tdatseqs, sdatseqs, smlseqs], [comouts, comouts]]
            else:
                return [[tdatseqs, sdatseqs, smlseqs], comouts]

    def divideseqs_threed(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        tdatseqs = list()
        sdatseqs = list()
        comouts = list()

        fiddat = dict()

        for fid in batchfids:
            wtdatseq = seqdata['dt%s' % (tt)][fid]
            wsdatseq = seqdata['ds%s' % (tt)][fid]
            
            wtdatseq = wtdatseq[:self.config['tdatlen']]

            newlen = self.config['sdatlen'] - len(wsdatseq)
            if newlen < 0:
                newlen=0
            wsdatseq = wsdatseq.tolist()
            for k in range(newlen):
                wsdatseq.append(np.zeros(self.config['stdatlen']))
            for i in range(0, len(wsdatseq)):
                wsdatseq[i] = np.array(wsdatseq[i])[:self.config['stdatlen']]
            wsdatseq = np.asarray(wsdatseq)
            wsdatseq = wsdatseq[:self.config['sdatlen'], :]
            
            if not self.training:
                fiddat[fid] = [wtdatseq, wsdatseq]
            else:
                tdatseqs.append(wtdatseq)
                sdatseqs.append(wsdatseq)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        tdatseqs = np.asarray(tdatseqs)
        sdatseqs = np.asarray(sdatseqs)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            if self.config['num_output'] == 2:
                return [[tdatseqs, sdatseqs], [comouts, comouts]]
            else:
                return [[tdatseqs, sdatseqs], comouts]

    def divideseqs_graphast(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        tdatseqs = list()
        smlnodes = list()
        smledges = list()
        comouts = list()

        fiddat = dict()

        for fid in batchfids:
            wtdatseq = seqdata['dt%s' % (tt)][fid]
            wsmlnodes = seqdata['s%s_nodes' % (tt)][fid]
            wsmledges = seqdata['s%s_edges' % (tt)][fid]
            if (wsmledges.shape[0]>1000):
                continue

            wsmlnodes = wsmlnodes[:self.config['maxastnodes']]
            tmp = np.zeros(self.config['maxastnodes'], dtype='int32')
            tmp[:wsmlnodes.shape[0]] = wsmlnodes
            wsmlnodes = np.int32(tmp)
            wsmledges = np.asarray(wsmledges.todense())
            wsmledges = wsmledges[:self.config['maxastnodes'], :self.config['maxastnodes']]
            tmp = np.zeros((self.config['maxastnodes'], self.config['maxastnodes']), dtype='int32')
            tmp[:wsmledges.shape[0], :wsmledges.shape[1]] = wsmledges
            wsmledges = np.int32(tmp)
            wtdatseq = wtdatseq[:self.config['tdatlen']]
            
            if not self.training:
                fiddat[fid] = [wtdatseq, wsmlnodes, wsmledges]
            else:
                if (self.config['use_tdats']):
                    tdatseqs.append(wtdatseq)
                smlnodes.append(wsmlnodes)
                smledges.append(wsmledges)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        if (self.config['use_tdats']):
            tdatseqs = np.asarray(tdatseqs)
        smlnodes = np.asarray(smlnodes)
        smledges = np.asarray(smledges)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            if self.config['num_output'] == 2:
                return [[tdatseqs, smlnodes, smledges], [comouts, comouts]]
            else:
                if (self.config['use_tdats']):
                    return [[tdatseqs, smlnodes, smledges], comouts]
                else:
                    return [[smlnodes, smledges], comouts]


    def divideseqs_graphast_threed(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        tdatseqs = list()
        sdatseqs = list()
        smlnodes = list()
        smledges = list()
        comouts = list()

        fiddat = dict()

        for fid in batchfids:
            wtdatseq = seqdata['dt%s' % (tt)][fid]
            wsdatseq = seqdata['ds%s' % (tt)][fid]
            wsmlnodes = seqdata['s%s_nodes' % (tt)][fid]
            wsmledges = seqdata['s%s_edges' % (tt)][fid]
            if (wsmledges.shape[0] > 1000):
                continue

            wsmlnodes = wsmlnodes[:self.config['maxastnodes']]
            tmp = np.zeros(self.config['maxastnodes'], dtype='int32')
            tmp[:wsmlnodes.shape[0]] = wsmlnodes
            wsmlnodes = np.int32(tmp)

            wsmledges = np.asarray(wsmledges.todense())
            wsmledges = wsmledges[:self.config['maxastnodes'], :self.config['maxastnodes']]
            tmp = np.zeros((self.config['maxastnodes'], self.config['maxastnodes']), dtype='int32')
            tmp[:wsmledges.shape[0], :wsmledges.shape[1]] = wsmledges
            wsmledges = np.int32(tmp)
            
            wtdatseq = wtdatseq[:self.config['tdatlen']]

            newlen = self.config['sdatlen'] - len(wsdatseq)
            if newlen < 0:
                newlen=0
            wsdatseq = wsdatseq.tolist()
            for k in range(newlen):
                wsdatseq.append(np.zeros(self.config['stdatlen']))
            for i in range(0, len(wsdatseq)):
                wsdatseq[i] = np.array(wsdatseq[i])[:self.config['stdatlen']]
            wsdatseq = np.asarray(wsdatseq)
            wsdatseq = wsdatseq[:self.config['sdatlen'], :]

            
            if not self.training:
                fiddat[fid] = [wtdatseq, wsdatseq, wsmlnodes, wsmledges]
            else:
                if (self.config['use_tdats']):
                    tdatseqs.append(wtdatseq)
                sdatseqs.append(wsdatseq)
                smlnodes.append(wsmlnodes)
                smledges.append(wsmledges)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        if (self.config['use_tdats']):
            tdatseqs = np.asarray(tdatseqs)
        sdatseqs = np.asarray(sdatseqs)
        smlnodes = np.asarray(smlnodes)
        smledges = np.asarray(smledges)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            if self.config['num_output'] == 2:
                return [[tdatseqs, sdatseqs, smlnodes, smledges], [comouts, comouts]]
            else:
                if (self.config['use_tdats']):
                    return [[tdatseqs, sdatseqs, smlnodes, smledges], comouts]
                else:
                    return [[sdatseqs, smlnodes, smledges], comouts]

    def idx2tok(self, nodelist, path):
        out = list()
        for idx in path:
            out.append(nodelist[idx])
        return out

    def divideseqs_pathast_threed(self, batchfids, seqdata, comvocabsize, tt):
        import keras.utils

        tdatseqs = list()
        sdatseqs = list()
        smlpaths = list()
        comouts = list()

        fiddat = dict()

        for fid in batchfids:
            wtdatseq = seqdata['dt%s' % (tt)][fid]
            wsdatseq = seqdata['ds%s' % (tt)][fid]
            wsmlnodes = seqdata['s%s_nodes' % (tt)][fid]
            wsmledges = seqdata['s%s_edges' % (tt)][fid]
            if (wsmledges.shape[0] > 1000):
                continue
            
            g = nx.from_numpy_matrix(wsmledges.todense())
            astpaths = nx.all_pairs_shortest_path(g, cutoff=self.config['pathlen'])
            wsmlpaths = list()

            for astpath in astpaths:
                source = astpath[0]
                
                if (len([n for n in g.neighbors(source)]) > 1):
                    continue
                
                for path in astpath[1].values():
                    if len([n for n in g.neighbors(path[-1])]) > 1:
                        continue
                    
                    if len(path) > 1 and len(path) <= self.config['pathlen']:
                        newpath = self.idx2tok(wsmlnodes, path)
                        tmp = [0] * (self.config['pathlen'] - len(newpath))
                        newpath.extend(tmp)
                        wsmlpaths.append(newpath)
            
            random.shuffle(wsmlpaths)
            wsmlpaths = wsmlpaths[:self.config['maxpaths']]
            if len(wsmlpaths) < self.config['maxpaths']:
                wsmlpaths.extend([[0]*self.config['pathlen']] * (self.config['maxpaths'] - len(wsmlpaths)))
            wsmlpaths = np.asarray(wsmlpaths)
            
            wtdatseq = wtdatseq[:self.config['tdatlen']]

            newlen = self.config['sdatlen'] - len(wsdatseq)
            if newlen < 0:
                newlen=0
            wsdatseq = wsdatseq.tolist()
            for k in range(newlen):
                wsdatseq.append(np.zeros(self.config['stdatlen']))
            for i in range(0, len(wsdatseq)):
                wsdatseq[i] = np.array(wsdatseq[i])[:self.config['stdatlen']]
            wsdatseq = np.asarray(wsdatseq)
            wsdatseq = wsdatseq[:self.config['sdatlen'], :]

            
            if not self.training:
                fiddat[fid] = [wtdatseq, wsdatseq, wsmlpaths]
            else:
                if (self.config['use_tdats']):
                    tdatseqs.append(wtdatseq)
                if (self.config['use_sdats']):
                    sdatseqs.append(wsdatseq)
                smlpaths.append(wsmlpaths)
                comout = keras.utils.to_categorical(self.firstwords[fid], num_classes=comvocabsize)
                comouts.append(np.asarray(comout))

        if (self.config['use_tdats']):
            tdatseqs = np.asarray(tdatseqs)
        if (self.config['use_sdats']):
            sdatseqs = np.asarray(sdatseqs)
        smlpaths = np.asarray(smlpaths)
        comouts = np.asarray(comouts)

        if not self.training:
            return fiddat
        else:
            if self.config['num_output'] == 2:
                return [[tdatseqs, sdatseqs, smlpaths], [comouts, comouts]]
            else:
                if (self.config['use_tdats'] and self.config['use_sdats']):
                    return [[tdatseqs, sdatseqs, smlpaths], comouts]
                elif (self.config['use_tdats'] and not self.config['use_sdats']):
                    return [[tdatseqs, smlpaths], comouts]
                elif (not self.config['use_tdats'] and self.config['use_sdats']):
                    return [[sdatseqs, smlpaths], comouts]
                else:
                    return [[smlpaths], comouts]