nn.cpp

#include "nn.hpp"
#include <algorithm>
#include <cmath>
#include <fstream>
#include <numeric>
#include <stdexcept>
#include <cassert>

float NeuralNet::xavier(int fan_in, int fan_out) {
    float limit = std::sqrt(6.0f / (fan_in + fan_out));
    std::uniform_real_distribution<float> dist(-limit, limit);
    return dist(rng_);
}

void NeuralNet::relu(std::vector<float>& v) {
    for (auto& x : v) x = x > 0.f ? x : 0.f;
}

void NeuralNet::softmax(std::vector<float>& v) {
    float mx = *std::max_element(v.begin(), v.end());
    float sum = 0.f;
    for (auto& x : v) { x = std::exp(x - mx); sum += x; }
    for (auto& x : v) x /= sum;
}



void NeuralNet::init(const std::vector<std::string>& class_labels, unsigned seed) {
    rng_.seed(seed);
    labels      = class_labels;
    num_classes = static_cast<int>(labels.size());

    embed.resize(VOCAB_SIZE * EMBED_DIM);
    for (auto& e : embed) e = xavier(VOCAB_SIZE, EMBED_DIM);

    w1.resize(HIDDEN1 * EMBED_DIM);
    b1.assign(HIDDEN1, 0.f);
    for (auto& w : w1) w = xavier(EMBED_DIM, HIDDEN1);

    w2.resize(HIDDEN2 * HIDDEN1);
    b2.assign(HIDDEN2, 0.f);
    for (auto& w : w2) w = xavier(HIDDEN1, HIDDEN2);

    wo.resize(num_classes * HIDDEN2);
    bo.assign(num_classes, 0.f);
    for (auto& w : wo) w = xavier(HIDDEN2, num_classes);

    a0.resize(EMBED_DIM);
    a1.resize(HIDDEN1);
    a2.resize(HIDDEN2);
    ao.resize(num_classes);
}


std::vector<float> NeuralNet::forward(const std::vector<int>& tokens) {
    std::fill(a0.begin(), a0.end(), 0.f);
    int n = 0;
    for (int tid : tokens) {
        if (tid < 0 || tid >= VOCAB_SIZE) continue;
        for (int d = 0; d < EMBED_DIM; ++d)
            a0[d] += embed[tid * EMBED_DIM + d];
        ++n;
    }
    if (n > 0) for (auto& x : a0) x /= n;

    std::vector<float> pre1(HIDDEN1, 0.f);
    for (int o = 0; o < HIDDEN1; ++o) {
        pre1[o] = b1[o];
        for (int i = 0; i < EMBED_DIM; ++i)
            pre1[o] += w1[o * EMBED_DIM + i] * a0[i];
    }
    a1 = pre1;
    relu(a1);

    std::vector<float> pre2(HIDDEN2, 0.f);
    for (int o = 0; o < HIDDEN2; ++o) {
        pre2[o] = b2[o];
        for (int i = 0; i < HIDDEN1; ++i)
            pre2[o] += w2[o * HIDDEN1 + i] * a1[i];
    }
    a2 = pre2;
    relu(a2);

    std::vector<float> preo(num_classes, 0.f);
    for (int o = 0; o < num_classes; ++o) {
        preo[o] = bo[o];
        for (int i = 0; i < HIDDEN2; ++i)
            preo[o] += wo[o * HIDDEN2 + i] * a2[i];
    }
    ao = preo;
    softmax(ao);

    return ao;
}

int NeuralNet::predict(const std::vector<int>& tokens) {
    auto probs = forward(tokens);
    return static_cast<int>(
        std::max_element(probs.begin(), probs.end()) - probs.begin());
}

std::string NeuralNet::predict_label(const std::vector<int>& tokens) {
    int idx = predict(tokens);
    return labels[idx];
}


float NeuralNet::train_step(const std::vector<int>& tokens, int label, float lr) {
    std::fill(a0.begin(), a0.end(), 0.f);
    int n = 0;
    for (int tid : tokens) {
        if (tid < 0 || tid >= VOCAB_SIZE) continue;
        for (int d = 0; d < EMBED_DIM; ++d)
            a0[d] += embed[tid * EMBED_DIM + d];
        ++n;
    }
    if (n > 0) for (auto& x : a0) x /= n;

    std::vector<float> pre1(HIDDEN1, 0.f);
    for (int o = 0; o < HIDDEN1; ++o) {
        pre1[o] = b1[o];
        for (int i = 0; i < EMBED_DIM; ++i)
            pre1[o] += w1[o * EMBED_DIM + i] * a0[i];
    }
    a1 = pre1; relu(a1);

    std::vector<float> pre2(HIDDEN2, 0.f);
    for (int o = 0; o < HIDDEN2; ++o) {
        pre2[o] = b2[o];
        for (int i = 0; i < HIDDEN1; ++i)
            pre2[o] += w2[o * HIDDEN1 + i] * a1[i];
    }
    a2 = pre2; relu(a2);

    std::vector<float> preo(num_classes, 0.f);
    for (int o = 0; o < num_classes; ++o) {
        preo[o] = bo[o];
        for (int i = 0; i < HIDDEN2; ++i)
            preo[o] += wo[o * HIDDEN2 + i] * a2[i];
    }
    ao = preo; softmax(ao);

    // Cross-entropy loss
    float loss = -std::log(std::max(ao[label], 1e-9f));

    std::vector<float> d_preo = ao;
    d_preo[label] -= 1.f;

    // Output
    std::vector<float> d_a2(HIDDEN2, 0.f);
    for (int o = 0; o < num_classes; ++o) {
        bo[o] -= lr * d_preo[o];
        for (int i = 0; i < HIDDEN2; ++i) {
            d_a2[i] += wo[o * HIDDEN2 + i] * d_preo[o];
            wo[o * HIDDEN2 + i] -= lr * d_preo[o] * a2[i];
        }
    }

    // Layer 2 Backprop ReLU
    std::vector<float> d_pre2(HIDDEN2, 0.f);
    for (int i = 0; i < HIDDEN2; ++i)
        d_pre2[i] = (pre2[i] > 0.f) ? d_a2[i] : 0.f;

    std::vector<float> d_a1(HIDDEN1, 0.f);
    for (int o = 0; o < HIDDEN2; ++o) {
        b2[o] -= lr * d_pre2[o];
        for (int i = 0; i < HIDDEN1; ++i) {
            d_a1[i] += w2[o * HIDDEN1 + i] * d_pre2[o];
            w2[o * HIDDEN1 + i] -= lr * d_pre2[o] * a1[i];
        }
    }

    // Layer 1 backprop ReLU
    std::vector<float> d_pre1(HIDDEN1, 0.f);
    for (int i = 0; i < HIDDEN1; ++i)
        d_pre1[i] = (pre1[i] > 0.f) ? d_a1[i] : 0.f;

    std::vector<float> d_a0(EMBED_DIM, 0.f);
    for (int o = 0; o < HIDDEN1; ++o) {
        b1[o] -= lr * d_pre1[o];
        for (int i = 0; i < EMBED_DIM; ++i) {
            d_a0[i] += w1[o * EMBED_DIM + i] * d_pre1[o];
            w1[o * EMBED_DIM + i] -= lr * d_pre1[o] * a0[i];
        }
    }

    if (n > 0) {
        float scale = 1.f / n;
        for (int tid : tokens) {
            if (tid < 0 || tid >= VOCAB_SIZE) continue;
            for (int d = 0; d < EMBED_DIM; ++d)
                embed[tid * EMBED_DIM + d] -= lr * d_a0[d] * scale;
        }
    }

    return loss;
}


bool NeuralNet::save(const std::string& path) const {
    std::ofstream f(path, std::ios::binary);
    if (!f) return false;

    uint32_t magic = 0x4E564953; // 'NVIS'
    f.write(reinterpret_cast<const char*>(&magic), 4);
    uint32_t nc = static_cast<uint32_t>(num_classes);
    f.write(reinterpret_cast<const char*>(&nc), 4);

    for (const auto& lbl : labels) {
        uint32_t len = static_cast<uint32_t>(lbl.size());
        f.write(reinterpret_cast<const char*>(&len), 4);
        f.write(lbl.data(), len);
    }

    auto write_vec = [&](const std::vector<float>& v) {
        f.write(reinterpret_cast<const char*>(v.data()),
                static_cast<std::streamsize>(v.size() * sizeof(float)));
    };
    write_vec(embed);
    write_vec(w1); write_vec(b1);
    write_vec(w2); write_vec(b2);
    write_vec(wo); write_vec(bo);

    return f.good();
}

bool NeuralNet::load(const std::string& path) {
    std::ifstream f(path, std::ios::binary);
    if (!f) return false;

    uint32_t magic = 0;
    f.read(reinterpret_cast<char*>(&magic), 4);
    if (magic != 0x4E564953) return false;

    uint32_t nc = 0;
    f.read(reinterpret_cast<char*>(&nc), 4);
    num_classes = static_cast<int>(nc);

    labels.resize(num_classes);
    for (auto& lbl : labels) {
        uint32_t len = 0;
        f.read(reinterpret_cast<char*>(&len), 4);
        lbl.resize(len);
        f.read(lbl.data(), len);
    }

    auto read_vec = [&](std::vector<float>& v, size_t n) {
        v.resize(n);
        f.read(reinterpret_cast<char*>(v.data()),
               static_cast<std::streamsize>(n * sizeof(float)));
    };
    read_vec(embed, VOCAB_SIZE * EMBED_DIM);
    read_vec(w1, HIDDEN1 * EMBED_DIM); read_vec(b1, HIDDEN1);
    read_vec(w2, HIDDEN2 * HIDDEN1);   read_vec(b2, HIDDEN2);
    read_vec(wo, num_classes * HIDDEN2); read_vec(bo, num_classes);

    a0.resize(EMBED_DIM);
    a1.resize(HIDDEN1);
    a2.resize(HIDDEN2);
    ao.resize(num_classes);

    return f.good();
}

int NeuralNet::param_count() const {
    return static_cast<int>(
        embed.size() + w1.size() + b1.size() +
        w2.size() + b2.size() +
        wo.size() + bo.size());
}