graph.cpp

// graph.cpp — parse + execute a minimal sequential inference graph.
//
// Python equivalent: onnxruntime.InferenceSession(path).run(None, {"x": x}),
// for a model that is just nn.Sequential(Linear, ReLU, Linear, Sigmoid, ...).
#include "cppml/graph.hpp"

#include <cmath>
#include <fstream>
#include <sstream>
#include <stdexcept>

namespace cppml {

const char* Node::kind_name() const {
    switch (kind) {
        case Kind::Linear:  return "linear";
        case Kind::ReLU:    return "relu";
        case Kind::Sigmoid: return "sigmoid";
    }
    return "?";
}

Tensor<float> Node::apply(const Tensor<float>& x) const {
    switch (kind) {
        case Kind::Linear: {
            // y = x @ W + b, bias broadcast over the batch dimension.
            Tensor<float> y = x.dot(weight);
            const std::size_t batch = y.shape()[0];
            for (std::size_t r = 0; r < batch; ++r)
                for (std::size_t c = 0; c < out_features; ++c) y(r, c) += bias(0, c);
            return y;
        }
        case Kind::ReLU:
            return x.map([](float v) { return v > 0.0f ? v : 0.0f; });
        case Kind::Sigmoid:
            return x.map([](float v) { return 1.0f / (1.0f + std::exp(-v)); });
    }
    throw std::logic_error("Node::apply: unknown kind");
}

namespace {
// Read exactly `count` floats from the token stream, erroring (with the line
// label) if the stream runs dry — weights may span several physical lines.
std::vector<float> read_floats(std::istream& in, std::size_t count, const char* what) {
    std::vector<float> v(count);
    for (std::size_t i = 0; i < count; ++i) {
        if (!(in >> v[i]))
            throw std::runtime_error(std::string("graph: ran out of values reading ") + what);
    }
    return v;
}
}  // namespace

Graph Graph::load(std::istream& in) {
    Graph g;
    bool header_seen = false, input_seen = false;
    std::string line;
    std::size_t lineno = 0;

    auto fail = [&](const std::string& msg) {
        throw std::runtime_error("graph: line " + std::to_string(lineno) + ": " + msg);
    };

    while (std::getline(in, line)) {
        ++lineno;
        // Strip comments (everything after '#') and skip blank lines.
        const auto hash = line.find('#');
        if (hash != std::string::npos) line.erase(hash);
        std::istringstream ls(line);
        std::string tok;
        if (!(ls >> tok)) continue;  // blank / comment-only

        if (!header_seen) {
            if (tok != "cppml-graph") fail("expected 'cppml-graph <version>' header");
            int version = 0;
            if (!(ls >> version) || version != 1) fail("unsupported graph version");
            header_seen = true;
            continue;
        }

        if (tok == "input") {
            if (!(ls >> g.input_features_)) fail("input requires a feature count");
            input_seen = true;
        } else if (tok == "linear") {
            if (!input_seen) fail("'input' must precede the first layer");
            Node n;
            n.kind = Node::Kind::Linear;
            if (!(ls >> n.in_features >> n.out_features))
                fail("linear requires <in_features> <out_features>");
            n.weight = Tensor<float>({n.in_features, n.out_features});
            n.bias = Tensor<float>({1, n.out_features}, 0.0f);
            g.nodes_.push_back(std::move(n));
        } else if (tok == "weights") {
            if (g.nodes_.empty() || g.nodes_.back().kind != Node::Kind::Linear)
                fail("'weights' must follow a 'linear' line");
            Node& n = g.nodes_.back();
            auto vals = read_floats(ls, n.in_features * n.out_features, "weights");
            for (std::size_t i = 0; i < vals.size(); ++i) n.weight[i] = vals[i];
        } else if (tok == "bias") {
            if (g.nodes_.empty() || g.nodes_.back().kind != Node::Kind::Linear)
                fail("'bias' must follow a 'linear' line");
            Node& n = g.nodes_.back();
            auto vals = read_floats(ls, n.out_features, "bias");
            for (std::size_t i = 0; i < vals.size(); ++i) n.bias[i] = vals[i];
        } else if (tok == "relu") {
            Node n; n.kind = Node::Kind::ReLU; g.nodes_.push_back(std::move(n));
        } else if (tok == "sigmoid") {
            Node n; n.kind = Node::Kind::Sigmoid; g.nodes_.push_back(std::move(n));
        } else {
            fail("unknown directive '" + tok + "'");
        }
    }

    if (!header_seen) throw std::runtime_error("graph: missing 'cppml-graph' header");
    if (!input_seen) throw std::runtime_error("graph: missing 'input' line");
    return g;
}

Graph Graph::load_file(const std::string& path) {
    std::ifstream f(path);
    if (!f) throw std::runtime_error("graph: cannot open file '" + path + "'");
    return load(f);
}

Tensor<float> Graph::forward(const Tensor<float>& x) const {
    if (x.ndim() != 2 || x.shape()[1] != input_features_)
        throw std::invalid_argument("graph: input must be (batch, input_features)");
    Tensor<float> cur = x;
    for (const auto& node : nodes_) cur = node.apply(cur);
    return cur;
}

std::string Graph::summary() const {
    std::ostringstream os;
    os << "Graph(input=" << input_features_ << ", " << nodes_.size() << " nodes)\n";
    for (std::size_t i = 0; i < nodes_.size(); ++i) {
        os << "  [" << i << "] " << nodes_[i].kind_name();
        if (nodes_[i].kind == Node::Kind::Linear)
            os << " " << nodes_[i].in_features << "->" << nodes_[i].out_features;
        os << "\n";
    }
    return os.str();
}

}  // namespace cppml