<commit_before>#bin/python
# Main functions/classes for bri-focal v2.

import numpy as np
# DO NOT PLAY NICE WITH NANS
np.seterr(all='raise')
np.seterr(under='warn')
from copy import deepcopy
import sys
import gzip
import time
import re
from copy import deepcopy
#import pathos.multiprocessing as mp

# --- CONSTANTS --- #
THEANO=False
if THEANO:
    print 'WARNING: Asking for theano functions, but they are not declared here.'
    from theano import function, shared, scan
    import theano.tensor as tten
    from bf2f_theano_params import *
# yolo
#linn = mp.ProcessingPool(5)

class data_stream(object):
    """
    Class for data stream.
    (can use this as a generator)
    """
    def __init__(self, path):
        self.path = path
    def __iter__(self):
        """
        Just spits out lines from the file.
        """
        fi = gzip.open(self.path,'r')
        header = fi.readline()
        while True:
            line = fi.readline()
            if len(line) == 0: # EOF
                break
            else:
                example = map(int, line.split())
                yield example
    def get_vocab_sizes(self):
        """
        The first line of the data file should contain W, R.
        """
        fi = gzip.open(self.path,'r')
        header = fi.readline()
        fi.close()
        values = map(int, header.split())
        if len(values) == 2:
            W, R = values
        else:
            sys.exit('ERROR: data file incorrectly formatted.')
        return W, R
    def acquire_all(self):
        """
        Just suck it all in!
        """
        traindata = [[0, 0, 0]]
        fi = gzip.open(self.path, 'r')
        header = fi.readline()
        for line in fi:
            s, r, t = map(int, line.split())
            traindata.append([s, r, t])
        return np.array(traindata[1:])

class params(object):
    """
    Parameter object.
    Contains C, G, V and velocities for all.
    """
    def __init__(self, initial_parameters, fix_words=False, fix_relas=False):
        C, G, V = initial_parameters
        if C.shape != V.shape:
            raise ValueError
        if G.shape[1] != C.shape[1]:
            raise ValueError
        if G.shape[2] != C.shape[1]:
            raise ValueError
        self.W = C.shape[0]
        self.R = G.shape[0]
        self.d = C.shape[1] - 1
        # weights
        self.C = deepcopy(C)
        self.G = deepcopy(G)
        self.V = deepcopy(V)
        # velocities
        self.C_vel = np.zeros(shape=self.C.shape)
        self.G_vel = np.zeros(shape=self.G.shape)
        self.V_vel = np.zeros(shape=self.V.shape)
        # fix some parameters?
        # (never update these)
        self.fix_words = fix_words
        self.fix_relas = fix_relas

    def update(self, delta_parameters, alpha, mu):
        """
        Updates velocities and then parameters.
        """
        # unwrap
        deltaC, deltaG, deltaV = delta_parameters
        alphaC, alphaG, alphaV = alpha
        muC, muG, muV = mu
        # update velocities
        if not self.fix_words:
            self.C_vel = muC*self.C_vel + (1-muC)*deltaC
            self.V_vel = muV*self.V_vel + (1-muV)*deltaV
        if not self.fix_relas:
            self.G_vel = muG*self.G_vel + (1-muG)*deltaG
        # update parameters
        if not self.fix_words:
            self.C += alphaC*self.C_vel
            self.V += alphaV*self.V_vel
        if not self.fix_relas:
            self.G += alphaG*self.G_vel

    def grad_E(self, locations):
        """
        Gradients of the energy, evaluated at a list of triples.
        NOTE: this clearly depends on the choice of energy.
        Returns tensors whose first index corresponds to the input triple list.
        """
        C_sub = self.C[locations[:, 0]]
        G_sub = self.G[locations[:, 1]]
        V_sub = self.V[locations[:, 2]]
        # this is for Etype == 'dot'
        # TODO: make this efficient
        dE_C = -np.einsum('...i,...ij', V_sub, G_sub)
        dE_G = -np.einsum('...i,...j', V_sub, C_sub)
        dE_V = -np.einsum('...ij,...j', G_sub, C_sub)
        return dE_C, dE_G, dE_V

    def E_axis(self, triple, switch):
        """
        Returns energies over an axis (S, R, T) given two of the triple.
        """
        s, r, t = triple
        if switch == 'C':
            # return over all S
            #GC = np.dot(self.C, self.G[r].T)
            #energy = -np.dot(GC, self.V[t])
            # note: above version is significantly slower than the below
            VG = np.dot(self.V[t], self.G[r])
            energy = -np.dot(self.C, VG)
        elif switch == 'G':
            # return over all R
            VG = np.dot(self.V[t], self.G)
            energy = -np.dot(VG, self.C[s])
        elif switch == 'V':
            #return over all T
            GC = np.dot(self.G[r], self.C[s])
            energy = -np.dot(self.V, GC)
        else:
            print 'ERROR: Cannot parse switch.'
            sys.exit()
        return energy

    def E_triple(self, triple):
        """
        The energy of a SINGLE triple.
        """
        return -np.dot(self.V[triple[2]], np.dot(self.G[triple[1]], self.C[triple[0]]))

    def E(self, locations):
        """
        Just plain old energy between triples.
        locations is an array of triples.
        Outputs a list (same length as 'locations') of energy of each triple.
        """
        #C_sub = self.C[locations[:, 0]]
        #G_sub = self.G[locations[:, 1]]
        #V_sub = self.V[locations[:, 2]]
        # this is for Etype == 'dot'
        # TODO:
        #   profile speed wrt order
        #   wrt just looping through locations
        #   # yolo
        # profiling...<commit_msg>now including default option<commit_after>#bin/python
# Main functions/classes for bri-focal v2.

import numpy as np
# DO NOT PLAY NICE WITH NANS
np.seterr(all='raise')
np.seterr(under='warn')
from copy import deepcopy
import sys
import gzip
import time
import re
from copy import deepcopy
#import pathos.multiprocessing as mp

# --- CONSTANTS --- #
THEANO=False
if THEANO:
    print 'WARNING: Asking for theano functions, but they are not declared here.'
    from theano import function, shared, scan
    import theano.tensor as tten
    from bf2f_theano_params import *
# yolo
#linn = mp.ProcessingPool(5)

class data_stream(object):
    """
    Class for data stream.
    (can use this as a generator)
    """
    def __init__(self, path):
        self.path = path
    def __iter__(self):
        """
        Just spits out lines from the file.
        """
        fi = gzip.open(self.path,'r')
        header = fi.readline()
        while True:
            line = fi.readline()
            if len(line) == 0: # EOF
                break
            else:
                example = map(int, line.split())
                yield example
    def get_vocab_sizes(self):
        """
        The first line of the data file should contain W, R.
        """
        fi = gzip.open(self.path,'r')
        header = fi.readline()
        fi.close()
        values = map(int, header.split())
        if len(values) == 2:
            W, R = values
        else:
            sys.exit('ERROR: data file incorrectly formatted.')
        return W, R
    def acquire_all(self):
        """
        Just suck it all in!
        """
        traindata = [[0, 0, 0]]
        fi = gzip.open(self.path, 'r')
        header = fi.readline()
        for line in fi:
            s, r, t = map(int, line.split())
            traindata.append([s, r, t])
        return np.array(traindata[1:])

class params(object):
    """
    Parameter object.
    Contains C, G, V and velocities for all.
    """
    def __init__(self, initial_parameters, fix_words=False, fix_relas=False):
        C, G, V = initial_parameters
        if C.shape != V.shape:
            raise ValueError
        if G.shape[1] != C.shape[1]:
            raise ValueError
        if G.shape[2] != C.shape[1]:
            raise ValueError
        self.W = C.shape[0]
        self.R = G.shape[0]
        self.d = C.shape[1] - 1
        # weights
        self.C = deepcopy(C)
        self.G = deepcopy(G)
        self.V = deepcopy(V)
        # velocities
        self.C_vel = np.zeros(shape=self.C.shape)
        self.G_vel = np.zeros(shape=self.G.shape)
        self.V_vel = np.zeros(shape=self.V.shape)
        # fix some parameters?
        # (never update these)
        self.fix_words = fix_words
        self.fix_relas = fix_relas

    def update(self, delta_parameters, alpha, mu):
        """
        Updates velocities and then parameters.
        """
        # unwrap
        deltaC, deltaG, deltaV = delta_parameters
        alphaC, alphaG, alphaV = alpha
        muC, muG, muV = mu
        # update velocities
        if not self.fix_words:
            self.C_vel = muC*self.C_vel + (1-muC)*deltaC
            self.V_vel = muV*self.V_vel + (1-muV)*deltaV
        if not self.fix_relas:
            self.G_vel = muG*self.G_vel + (1-muG)*deltaG
        # update parameters
        if not self.fix_words:
            self.C += alphaC*self.C_vel
            self.V += alphaV*self.V_vel
        if not self.fix_relas:
            self.G += alphaG*self.G_vel

    def grad_E(self, locations):
        """
        Gradients of the energy, evaluated at a list of triples.
        NOTE: this clearly depends on the choice of energy.
        Returns tensors whose first index corresponds to the input triple list.
        """
        C_sub = self.C[locations[:, 0]]
        G_sub = self.G[locations[:, 1]]
        V_sub = self.V[locations[:, 2]]
        # this is for Etype == 'dot'
        # TODO: make this efficient
        dE_C = -np.einsum('...i,...ij', V_sub, G_sub)
        dE_G = -np.einsum('...i,...j', V_sub, C_sub)
        dE_V = -np.einsum('...ij,...j', G_sub, C_sub)
        return dE_C, dE_G, dE_V

    def E_axis(self, triple, switch):
        """
        Returns energies over an axis (S, R, T) given two of the triple.
        """
        s, r, t = triple
        if switch == 'C':
            # return over all S
            #GC = np.dot(self.C, self.G[r].T)
            #energy = -np.dot(GC, self.V[t])
            # note: above version is significantly slower than the below
            VG = np.dot(self.V[t], self.G[r])
            energy = -np.dot(self.C, VG)
        elif switch == 'G':
            # return over all R
            VG = np.dot(self.V[t], self.G)
            energy = -np.dot(VG, self.C[s])
        elif switch == 'V':
            #return over all T
            GC = np.dot(self.G[r], self.C[s])
            energy = -np.dot(self.V, GC)
        else:
            print 'ERROR: Cannot parse switch.'
            sys.exit()
        return energy

    def E_triple(self, triple):
        """
        The energy of a SINGLE triple.
        """
        return -np.dot(self.V[triple[2]], np.dot(self.G[triple[1]], self.C[triple[0]]))

    def E(self, locations=None):
        """
        Just plain old energy between triples.
        locations is an array of triples.
        Outputs a list (same length as 'locations') of energy of each triple.
        """
        #C_sub = self.C[locations[:, 0]]
        #G_sub = self.G[locations[:, 1]]
        #V_sub = self.V[locations[:, 2]]
        # this is for Etype == 'dot'
        # TODO:
        #   profile speed wrt order
        #   wrt just looping through locations
        #   # yolo
        # profiling...
{'fraction_non_alphanumeric': 0.12347081440055925, 'fraction_numerical': 0.005068157986717931, 'mean_word_length': 1.7328080229226361, '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': '1004066', 'n_tokens_mistral': 3744, 'n_tokens_neox': 3681, 'n_words': 1042}