"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.createSimpleNN = void 0;
var tfjs_node_1 = require("@tensorflow/tfjs-node");
var constants_1 = require("./constants");
function createSimpleNN(maxLength, numClasses, embeddingsMatrix, params) {
    var _a, _b, _c;
    if (params === void 0) { params = {
        dropout: constants_1.DEFAULT_DROPOUT,
        units: constants_1.DEFAULT_UNITS,
        learningRate: constants_1.DEFAULT_LEARNING_RATE,
    }; }
    var inputs = tfjs_node_1.input({ name: 'InputLayer', shape: [maxLength] });
    var embeddingsLayer = tfjs_node_1.layers.embedding({
        name: 'EmbeddingsLayer',
        inputDim: embeddingsMatrix.shape[0],
        outputDim: embeddingsMatrix.shape[1],
        inputLength: maxLength,
        weights: [embeddingsMatrix],
        trainable: true,
    });
    var lstmLayer = tfjs_node_1.layers.lstm({
        name: 'LSTMLayer',
        units: (_a = params.units) !== null && _a !== void 0 ? _a : constants_1.DEFAULT_UNITS,
        dropout: (_b = params.dropout) !== null && _b !== void 0 ? _b : constants_1.DEFAULT_DROPOUT,
        recurrentDropout: 0.3,
    });
    var denseLayer = tfjs_node_1.layers.dense({
        name: 'DenseLayer',
        units: numClasses,
        activation: 'softmax',
    });
    var outputs = denseLayer.apply(lstmLayer.apply(embeddingsLayer.apply(inputs)));
    var textClassifierModel = tfjs_node_1.model({
        name: 'SimpleTextClassifier',
        inputs: inputs,
        outputs: outputs,
    });
    textClassifierModel.compile({
        optimizer: tfjs_node_1.train.adam((_c = params.learningRate) !== null && _c !== void 0 ? _c : constants_1.DEFAULT_LEARNING_RATE),
        loss: 'categoricalCrossentropy',
        metrics: ['accuracy'],
    });
    return textClassifierModel;
}
exports.createSimpleNN = createSimpleNN;
//# sourceMappingURL=simple-nn.js.map