"use strict";
// -*- mode: js; indent-tabs-mode: nil; js-basic-offset: 4 -*-
//
// This file is part of Genie
//
// Copyright 2020 The Board of Trustees of the Leland Stanford Junior University
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Author: Giovanni Campagna <gcampagn@cs.stanford.edu>
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 (k !== "default" && Object.prototype.hasOwnProperty.call(mod, k)) __createBinding(result, mod, k);
    __setModuleDefault(result, mod);
    return result;
};
var __importDefault = (this && this.__importDefault) || function (mod) {
    return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
const path_1 = __importDefault(require("path"));
const fs = __importStar(require("fs"));
const util_1 = __importDefault(require("util"));
const base_training_job_1 = __importDefault(require("./base_training_job"));
const exec_utils_1 = require("./exec-utils");
const DEFAULT_TRAINING_CONFIG = {
    task_name: 'almond',
    train_iterations: 100000,
    save_every: 2000,
    log_every: 500,
    val_every: 1000
};
class GenieNLPTrainingJob extends base_training_job_1.default {
    constructor(options) {
        super(options);
        this._config = {};
        Object.assign(this._config, DEFAULT_TRAINING_CONFIG);
        if (options.config)
            Object.assign(this._config, options.config);
    }
    get config() {
        return this._config;
    }
    async _train() {
        await util_1.default.promisify(fs.symlink)(path_1.default.resolve(this.datadir), path_1.default.resolve(this.workdir, this._config.task_name));
        const args = [
            'train',
            '--train_tasks', this._config.task_name,
            '--save', path_1.default.resolve(this.workdir, 'model'),
            '--cache', path_1.default.resolve(this.workdir, 'cache'),
            '--data', this.workdir,
            '--preserve_case',
            '--no_commit',
        ];
        if (process.env.GENIENLP_EMBEDDINGS)
            args.push('--embeddings', process.env.GENIENLP_EMBEDDINGS);
        else
            args.push('--embeddings', path_1.default.resolve(this.workdir, 'embeddings'));
        for (let key in this._config) {
            if (['thingpedia_snapshot', 'thingpedia_developer_key', 'synthetic_depth', 'task_name'].indexOf(key) >= 0)
                continue;
            if (typeof this._config[key] === 'boolean') {
                if (this._config[key])
                    args.push('--' + key);
            }
            else {
                args.push('--' + key, String(this._config[key]));
            }
        }
        this.metrics = {};
        await exec_utils_1.execCommand(this, 'genienlp', args, (line) => {
            // the line we are looking for has the form:
            // ...:train_contextual_almond:70000/100000:val_deca:...
            // or
            // ...:train_almond:70000/100000:val_deca:...
            const match = /train_[a-z_]+:([0-9]+)\/[0-9]+:val_[a-z_]+:(.*)$/.exec(line);
            if (match === null)
                return;
            this.progress = parseFloat(match[1]) / this._config.train_iterations;
            for (let metric of match[2].split(':')) {
                let [key, value] = metric.split('_');
                this.metrics[key] = parseFloat(value);
            }
        });
        if (this._killed)
            return;
        if (this.outputdir)
            await this._findBestModel();
    }
    async _findBestModel() {
        if (this._killed)
            return;
        const args = ['export',
            '--path', path_1.default.resolve(this.workdir, 'model'),
            '--output', path_1.default.resolve(this.outputdir)
        ];
        if (process.env.GENIENLP_EMBEDDINGS)
            args.push('--embeddings', process.env.GENIENLP_EMBEDDINGS);
        else
            args.push('--embeddings', path_1.default.resolve(this.workdir, 'embeddings'));
        await exec_utils_1.execCommand(this, 'genienlp', args);
    }
    async train() {
        if (this.outputdir) {
            // remove the output directory first
            // this will complain loudly if the directory is not empty
            await exec_utils_1.safeRmdir(this.outputdir);
            await exec_utils_1.safeMkdir(this.outputdir);
        }
        await exec_utils_1.safeMkdir(this.workdir);
        if (this._killed)
            return;
        await this._train();
    }
}
exports.default = GenieNLPTrainingJob;
//# sourceMappingURL=genienlp.js.map