"use strict";
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    Object.defineProperty(o, k2, { enumerable: true, get: function() { return m[k]; } });
}) : (function(o, m, k, k2) {
    if (k2 === undefined) k2 = k;
    o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
    Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
    o["default"] = v;
});
var __importStar = (this && this.__importStar) || function (mod) {
    if (mod && mod.__esModule) return mod;
    var result = {};
    if (mod != null) for (var k in mod) if (Object.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.TextEmbedding = void 0;
const model_interface_1 = require("./model_interface");
const UniversalSentenceEncoder = __importStar(require("@tensorflow-models/universal-sentence-encoder"));
let model;
let tokenizer;
/**
 * Text Embedding with Tensorflow Universal Sentence Encoder (USE)
 * @class TextEmbedding
 * @implements {TensorScriptModelInterface}
 */
class TextEmbedding extends model_interface_1.TensorScriptModelInterface {
    /**
     * @param {Object} options - Options for USE
     * @param {{model:Object,tf:Object,}} properties - extra instance properties
     */
    constructor(options = {}, properties) {
        const config = Object.assign({}, options);
        super(config, properties);
        this.type = 'TextEmbedding';
        return this;
    }
    /**
     * Asynchronously loads Universal Sentence Encoder and tokenizer
     * @override
     * @return {Object} returns loaded UniversalSentenceEncoder model
     */
    async train() {
        const promises = [];
        if (!model)
            promises.push(UniversalSentenceEncoder.load());
        else
            promises.push(Promise.resolve(model));
        if (!tokenizer)
            promises.push(UniversalSentenceEncoder.loadTokenizer());
        else
            promises.push(Promise.resolve(tokenizer));
        const USE = await Promise.all(promises);
        if (!model)
            model = USE[0];
        if (!tokenizer)
            tokenizer = USE[1];
        this.model = model;
        this.tokenizer = tokenizer;
        this.trained = true;
        this.compiled = true;
        return this.model;
    }
    /**
     * Calculates sentence embeddings
     * @override
     * @param {Array<Array<number>>|Array<number>} input_array - new test independent variables
     * @param {Object} options - model prediction options
     * @return {{data: Promise}} returns tensorflow prediction
     */
    calculate(input_array, options = {}) {
        if (!input_array || Array.isArray(input_array) === false)
            throw new Error('invalid input array of sentences');
        const embeddings = this.model.embed(input_array);
        return embeddings;
    }
    /**
     * Returns prediction values from tensorflow model
     * @param {Array<string>} input_matrix - array of sentences to embed
     * @param {Boolean} [options.json=true] - return object instead of typed array
     * @param {Boolean} [options.probability=true] - return real values instead of integers
     * @return {Array<Array<number>>} predicted model values
     */
    async predict(input_array, options = {}) {
        const config = Object.assign({
            json: true,
            probability: true,
        }, options);
        const embeddings = await this.calculate(input_array, options);
        const predictions = await embeddings.data();
        if (config.json === false) {
            return predictions;
        }
        else {
            const shape = [input_array.length, 512,];
            const predictionValues = (options.probability === false)
                ? Array.from(predictions).map(Math.round)
                : Array.from(predictions);
            return this.reshape(predictionValues, shape);
        }
    }
}
exports.TextEmbedding = TextEmbedding;
