"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.createBiLstmModel = void 0;
var tfjs_node_1 = require("@tensorflow/tfjs-node");
var DEFAULT_DROPOUT = 0.1;
var DEFAULT_UNITS = 128;
var DEFAULT_LEARNING_RATE = 0.001;
function createBiLstmModel(maxLength, entities, embeddingsMatrix, params) {
    var _a, _b, _c;
    if (params === void 0) { params = {
        dropout: DEFAULT_DROPOUT,
        units: DEFAULT_UNITS,
        learningRate: 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 dropoutLayer = tfjs_node_1.layers.dropout({
        name: 'DropoutLayer',
        rate: (_a = params.dropout) !== null && _a !== void 0 ? _a : DEFAULT_DROPOUT,
    });
    var bidirectionalLSTM = tfjs_node_1.layers.bidirectional({
        name: 'BidirectionalLayer',
        layer: tfjs_node_1.layers.lstm({
            units: (_b = params.units) !== null && _b !== void 0 ? _b : DEFAULT_UNITS,
            returnSequences: true,
            recurrentDropout: 0.1,
        }),
    });
    var timeDistributedLayer = tfjs_node_1.layers.timeDistributed({
        name: 'TimeDistributedLayer',
        layer: tfjs_node_1.layers.dense({
            units: entities.length,
            activation: 'softmax',
        }),
    });
    var outputs = timeDistributedLayer.apply(bidirectionalLSTM.apply(dropoutLayer.apply(embeddingsLayer.apply(inputs))));
    var nerModel = tfjs_node_1.model({ name: 'BiLstmNerModel', inputs: inputs, outputs: outputs });
    nerModel.compile({
        optimizer: tfjs_node_1.train.adam((_c = params.learningRate) !== null && _c !== void 0 ? _c : DEFAULT_LEARNING_RATE),
        loss: 'categoricalCrossentropy',
        metrics: ['acc'],
    });
    return nerModel;
}
exports.createBiLstmModel = createBiLstmModel;
//# sourceMappingURL=bilstm-model.js.map