<commit_before># cython: profile=True
# cython: cdivision=True
# cython: infer_types=True
from __future__ import print_function

from libc.string cimport memmove, memset, memcpy

cimport cython

from cymem.cymem cimport Pool
from preshed.maps cimport map_init as Map_init
from preshed.maps cimport map_set as Map_set
from preshed.maps cimport map_get as Map_get
from preshed.maps cimport map_iter as Map_iter
from preshed.maps cimport key_t

from .typedefs cimport weight_t, atom_t, feat_t
from .typedefs cimport len_t, idx_t
from .blas cimport MatMat, MatVec, VecVec, Vec
from .structs cimport MapC
from .structs cimport NeuralNetC
from .structs cimport ExampleC
from .structs cimport FeatureC
from .structs cimport EmbedC
from .structs cimport ConstantsC
from .structs cimport do_update_t

from .eg cimport Example

from .lvl0 cimport expf
from .lvl0 cimport adam
from .lvl0 cimport adadelta
from .lvl0 cimport adagrad
from .lvl0 cimport vanilla_sgd_update_step
from .lvl0 cimport dot_plus__ELU
from .lvl0 cimport d_ELU__dot
from .lvl0 cimport softmax
from .lvl0 cimport d_log_loss

import numpy


DEF USE_BATCH_NORM = False


cdef class NN:
    @staticmethod
    cdef int nr_weight(int nr_out, int nr_in) nogil:
        if USE_BATCH_NORM:
            return nr_out * nr_in + nr_out * 2
        else:
            return nr_out * nr_in + nr_out

    @staticmethod
    cdef void init(
        NeuralNetC* nn,
        Pool mem,
            widths,
            embed=None,
            update_step='adam',
            float eta=0.005,
            float eps=1e-6,
            float mu=0.2,
            float rho=1e-4,
            float bias=0.0,
            float alpha=0.0
    ) except *:
        if update_step == 'sgd':
            nn.update = vanilla_sgd_update_step
        elif update_step == 'adadelta':
            nn.update = adadelta
        elif update_step == 'adagrad':
            nn.update = adagrad
        else:
            nn.update = adam
        nn.feed_fwd = dot_plus__ELU
        nn.feed_bwd = d_ELU__dot

        nn.hp.t = 0
        nn.hp.a = alpha
        nn.hp.b = bias
        nn.hp.r = rho
        nn.hp.m = mu
        nn.hp.e = eta

        nn.nr_layer = len(widths)
        nn.widths = <len_t*>mem.alloc(nn.nr_layer, sizeof(widths[0]))
        cdef int i
        for i, width in enumerate(widths):
            nn.widths[i] = width

        nn.nr_weight = 0
        nn.nr_node = 0
        for i in range(nn.nr_layer-1):
            nn.nr_weight += NN.nr_weight(nn.widths[i+1], nn.widths[i])
            nn.nr_node += nn.widths[i]
        nn.weights = <float*>mem.alloc(nn.nr_weight, sizeof(nn.weights[0]))
        nn.gradient = <float*>mem.alloc(nn.nr_weight, sizeof(nn.weights[0]))
        nn.momentum = <float*>mem.alloc(nn.nr_weight * 2, sizeof(nn.weights[0]))
        nn.averages = <float*>mem.alloc(nn.nr_node * 4, sizeof(nn.weights[0]))
        
        if embed is not None:
            vector_widths, features = embed
            Embedding.init(&nn.embed, mem, vector_widths, features)

        W = nn.weights
        for i in range(nn.nr_layer-2):
            he_uniform_initializer(W,
                nn.widths[i+1] * nn.widths[i])
            W += NN.nr_weight(nn.widths[i+1], nn.widths[i])
    
    @staticmethod
    cdef void train_example(NeuralNetC* nn, Pool mem, ExampleC* eg) except *:
        nn.hp.t += 1
        Embedding.insert_missing(mem, &nn.embed,
            eg.features, eg.nr_feat)
        Embedding.set_input(eg.fwd_state[0],
            eg.features, eg.nr_feat, &nn.embed)
        NN.forward(eg.scores, eg.fwd_state,
            nn)
        NN.backward(eg.bwd_state, nn.gradient,
            eg.fwd_state, eg.costs, nn)
        nn.update(nn.weights, nn.momentum, nn.gradient,
            nn.nr_weight, &nn.hp)
        if eg.nr_feat != 0:
            Embedding.fine_tune(&nn.embed, nn.gradient,
                eg.bwd_state[0], nn.widths[0], eg.features, eg.nr_feat,
                &nn.hp, nn.update)
    
    @staticmethod
    cdef void forward(float* scores, float** fwd, const NeuralNetC* nn) nogil:
        cdef const float* W = nn.weights
        for i in range(nn.nr_layer-1):
            nn.feed_fwd(&fwd[i],
                W, &nn.widths[i], nn.nr_layer-(i+1))
            W += NN.nr_weight(nn.widths[i+1], nn.widths[i])
        memcpy(scores,
            fwd[nn.nr_layer-1], sizeof(scores[0]) * nn.widths[nn.nr_layer-1])

    @staticmethod
    cdef void backward(float** bwd, float* gradient,
            const float* const* fwd, const float* costs, const NeuralNetC* nn) nogil:
        d_log_loss(bwd[nn.nr_layer-1],
            costs, fwd[nn.nr_layer-1], nn.widths[nn.nr_layer-1])
        cdef const float* W = nn.weights + nn.nr_weight
        cdef float* G = gradient + nn.nr_weight
        for i in range(nn.nr_layer-2, 0, -1):
            W -= NN.nr_weight(nn.widths[i+1], nn.widths[i])
            G -= NN.nr_weight(nn.widths[i+1], nn.widths[i])
            nn.feed_bwd(G, &bwd[i],
                W, &fwd[i], &nn.widths[i], i)
        MatVec.T_dot(bwd[0], nn.weights, bwd[1], nn.widths[1], nn.widths[0])
        MatMat.add_outer_i(gradient,
            bwd[1], fwd[0], nn.widths[1], nn.widths[0])
        VecVec.add_i(gradient + (nn.widths[1] * nn.widths[0]),
            bwd[1], 1.0, nn.widths[1])


cdef class Embedding:
    cdef Pool mem
    cdef EmbedC* c

    @staticmethod
    cdef void init(EmbedC* self, Pool mem, vector_widths, features) except *: 
        assert max(features) < len(vector_widths)
        # Create tables, which may be shared between different features
        # e.g., we might have a feature for this word, and a feature for next
        # word. These occupy different parts of the input vector, but draw
        # from the same embedding table.
        uniq_weights = <MapC*>mem.alloc(len(vector_widths), sizeof(MapC))
        uniq_momentum = <MapC*>mem.alloc(len(vector_widths), sizeof(MapC))
        for i, width in enumerate(vector_widths):
            Map_init(mem, &uniq_weights[i], 8)<commit_msg>* Move step 0 of gradient calculation inside feed_bwd function.<commit_after># cython: profile=True
# cython: cdivision=True
# cython: infer_types=True
from __future__ import print_function

from libc.string cimport memmove, memset, memcpy

cimport cython

from cymem.cymem cimport Pool
from preshed.maps cimport map_init as Map_init
from preshed.maps cimport map_set as Map_set
from preshed.maps cimport map_get as Map_get
from preshed.maps cimport map_iter as Map_iter
from preshed.maps cimport key_t

from .typedefs cimport weight_t, atom_t, feat_t
from .typedefs cimport len_t, idx_t
from .blas cimport MatMat, MatVec, VecVec, Vec
from .structs cimport MapC
from .structs cimport NeuralNetC
from .structs cimport ExampleC
from .structs cimport FeatureC
from .structs cimport EmbedC
from .structs cimport ConstantsC
from .structs cimport do_update_t

from .eg cimport Example

from .lvl0 cimport expf
from .lvl0 cimport adam
from .lvl0 cimport adadelta
from .lvl0 cimport adagrad
from .lvl0 cimport vanilla_sgd_update_step
from .lvl0 cimport dot_plus__ELU
from .lvl0 cimport d_ELU__dot
from .lvl0 cimport softmax
from .lvl0 cimport d_log_loss

import numpy


DEF USE_BATCH_NORM = False


cdef class NN:
    @staticmethod
    cdef int nr_weight(int nr_out, int nr_in) nogil:
        if USE_BATCH_NORM:
            return nr_out * nr_in + nr_out * 2
        else:
            return nr_out * nr_in + nr_out

    @staticmethod
    cdef void init(
        NeuralNetC* nn,
        Pool mem,
            widths,
            embed=None,
            update_step='adam',
            float eta=0.005,
            float eps=1e-6,
            float mu=0.2,
            float rho=1e-4,
            float bias=0.0,
            float alpha=0.0
    ) except *:
        if update_step == 'sgd':
            nn.update = vanilla_sgd_update_step
        elif update_step == 'adadelta':
            nn.update = adadelta
        elif update_step == 'adagrad':
            nn.update = adagrad
        else:
            nn.update = adam
        nn.feed_fwd = dot_plus__ELU
        nn.feed_bwd = d_ELU__dot

        nn.hp.t = 0
        nn.hp.a = alpha
        nn.hp.b = bias
        nn.hp.r = rho
        nn.hp.m = mu
        nn.hp.e = eta

        nn.nr_layer = len(widths)
        nn.widths = <len_t*>mem.alloc(nn.nr_layer, sizeof(widths[0]))
        cdef int i
        for i, width in enumerate(widths):
            nn.widths[i] = width

        nn.nr_weight = 0
        nn.nr_node = 0
        for i in range(nn.nr_layer-1):
            nn.nr_weight += NN.nr_weight(nn.widths[i+1], nn.widths[i])
            nn.nr_node += nn.widths[i]
        nn.weights = <float*>mem.alloc(nn.nr_weight, sizeof(nn.weights[0]))
        nn.gradient = <float*>mem.alloc(nn.nr_weight, sizeof(nn.weights[0]))
        nn.momentum = <float*>mem.alloc(nn.nr_weight * 2, sizeof(nn.weights[0]))
        nn.averages = <float*>mem.alloc(nn.nr_node * 4, sizeof(nn.weights[0]))
        
        if embed is not None:
            vector_widths, features = embed
            Embedding.init(&nn.embed, mem, vector_widths, features)

        W = nn.weights
        for i in range(nn.nr_layer-2):
            he_uniform_initializer(W,
                nn.widths[i+1] * nn.widths[i])
            W += NN.nr_weight(nn.widths[i+1], nn.widths[i])
    
    @staticmethod
    cdef void train_example(NeuralNetC* nn, Pool mem, ExampleC* eg) except *:
        nn.hp.t += 1
        Embedding.insert_missing(mem, &nn.embed,
            eg.features, eg.nr_feat)
        Embedding.set_input(eg.fwd_state[0],
            eg.features, eg.nr_feat, &nn.embed)
        NN.forward(eg.scores, eg.fwd_state,
            nn)
        NN.backward(eg.bwd_state, nn.gradient,
            eg.fwd_state, eg.costs, nn)
        nn.update(nn.weights, nn.momentum, nn.gradient,
            nn.nr_weight, &nn.hp)
        if eg.nr_feat != 0:
            Embedding.fine_tune(&nn.embed, nn.gradient,
                eg.bwd_state[0], nn.widths[0], eg.features, eg.nr_feat,
                &nn.hp, nn.update)
    
    @staticmethod
    cdef void forward(float* scores, float** fwd, const NeuralNetC* nn) nogil:
        cdef const float* W = nn.weights
        for i in range(nn.nr_layer-1):
            nn.feed_fwd(&fwd[i],
                W, &nn.widths[i], nn.nr_layer-(i+1))
            W += NN.nr_weight(nn.widths[i+1], nn.widths[i])
        memcpy(scores,
            fwd[nn.nr_layer-1], sizeof(scores[0]) * nn.widths[nn.nr_layer-1])

    @staticmethod
    cdef void backward(float** bwd, float* gradient,
            const float* const* fwd, const float* costs, const NeuralNetC* nn) nogil:
        d_log_loss(bwd[nn.nr_layer-1],
            costs, fwd[nn.nr_layer-1], nn.widths[nn.nr_layer-1])
        cdef const float* W = nn.weights + nn.nr_weight
        cdef float* G = gradient + nn.nr_weight
        for i in range(nn.nr_layer-2, -1, -1):
            W -= NN.nr_weight(nn.widths[i+1], nn.widths[i])
            G -= NN.nr_weight(nn.widths[i+1], nn.widths[i])
            nn.feed_bwd(G, &bwd[i],
                W, &fwd[i], &nn.widths[i], i)


cdef class Embedding:
    cdef Pool mem
    cdef EmbedC* c

    @staticmethod
    cdef void init(EmbedC* self, Pool mem, vector_widths, features) except *: 
        assert max(features) < len(vector_widths)
        # Create tables, which may be shared between different features
        # e.g., we might have a feature for this word, and a feature for next
        # word. These occupy different parts of the input vector, but draw
        # from the same embedding table.
        uniq_weights = <MapC*>mem.alloc(len(vector_widths), sizeof(MapC))
        uniq_momentum = <MapC*>mem.alloc(len(vector_widths), sizeof(MapC))
        for i, width in enumerate(vector_widths):
            Map_init(mem, &uniq_weights[i], 8)
{'fraction_non_alphanumeric': 0.10981923109564627, 'fraction_numerical': 0.011287448018331495, 'mean_word_length': 2.2606530160486993, 'pattern_counts': {'":': 0, '<': 19, '<?xml version=': 0, '>': 17, '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': '597987', 'n_tokens_mistral': 4280, 'n_tokens_neox': 4135, 'n_words': 1004}