text: /**
   * @license
   * Copyright 2019 Google LLC. All Rights Reserved.
   * 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.
   * =============================================================================
   */
  
  import * as tf from '@tensorflow/tfjs';
  import {ArgumentParser} from 'argparse';
  import * as fs from 'fs';
  import * as path from 'path';
  import * as shelljs from 'shelljs';
  
  import {loadData, loadMetadataTemplate} from './data';
  import {writeEmbeddingMatrixAndLabels} from './embedding';
  import { parse } from 'path';
  
  /**
   * Create a model for IMDB sentiment analysis.
   *
   * @param {string} modelType Type of the model to be created.
   * @param {number} vocabularySize Input vocabulary size.
   * @param {number} embeddingSize Embedding vector size, used to
   *   configure the embedding layer.
   * @returns An uncompiled instance of `tf.Model`.
   */
  export function buildModel(modelType, maxLen, vocabularySize, embeddingSize,
    numUnits, returnSeq) {
    // TODO(cais): Bidirectional and dense-only.
    const model = tf.sequential();
    if (modelType === 'multihot') {
      // A 'multihot' model takes a multi-hot encoding of all words in the
      // sentence and uses dense layers with relu and sigmoid activation functions
      // to classify the sentence.
      model.add(tf.layers.dense({
        units: 16,
        activation: 'relu',
        inputShape: [vocabularySize]
      }));
      model.add(tf.layers.dense({
        units: 16,
        activation: 'relu'
      }));
    } else {
      // All other model types use word embedding.
      model.add(tf.layers.embedding({
        inputDim: vocabularySize,
        outputDim: embeddingSize,
        inputLength: maxLen
      }));
      if (modelType === 'flatten') {
        model.add(tf.layers.flatten());
      } else if (modelType === 'cnn') {
        model.add(tf.layers.dropout({rate: 0.5}));
        model.add(tf.layers.conv1d({
          filters: 250,
          kernelSize: 5,
          strides: 1,
          padding: 'valid',
          activation: 'relu'
        }));
        model.add(tf.layers.globalMaxPool1d({}));
        model.add(tf.layers.dense({units: 250, activation: 'relu'}));
      } else if (modelType === 'simpleRNN') {
        model.add(tf.layers.simpleRNN({units: numUnits}));
      } else if (modelType === 'lstm') {
        // the shape of lstm output with return sequences cannot be accepted by dense layer,
        // it can work to add a flatten layer, but the score doesn't make sence
        if (returnSeq === 'true') {
          model.add(tf.layers.lstm({units: numUnits, returnSequences: returnSeq}));
          model.add(tf.layers.flatten())
        } else {
          model.add(tf.layers.lstm({units: numUnits}));
        }
      } else if (modelType === 'bidirectionalLSTM') {
        model.add(tf.layers.bidirectional(
            {layer: tf.layers.lstm({units: 32}), mergeMode: 'concat'}));
      } else {
        throw new Error(`Unsupported model type: ${modelType}`);
      }
    }
    if (returnSeq === 'true'){
      model.add(tf.layers.dense({units:maxLen, activation: 'softmax'}));
    } else {
      model.add(tf.layers.dense({units: 1, activation: 'sigmoid'}));
    }
    return model;
  }
  
  function parseArguments() {
    const parser = new ArgumentParser(
        {description: 'Train a model for IMDB sentiment analysis'});
    parser.addArgument('modelType', {
      type: 'string',
      optionStrings: [
         'multihot', 'flatten', 'cnn', 'simpleRNN', 'lstm', 'bidirectionalLSTM'],
      help: 'Model type'
    });
    parser.addArgument('--numWords', {
      type: 'int',
      defaultValue: 10000,
      help: 'Number of words in the vocabulary'
    });
    parser.addArgument('--maxLen', {
      type: 'int',
      defaultValue: 100,
      help: 'Maximum sentence length in number of words. ' +
          'Shorter sentences will be padded; longers ones will be truncated.'
    });
    parser.addArgument('--embeddingSize', {
      type: 'int',
      defaultValue: 128,
      help: 'Number of word embedding dimensions'
    });
    parser.addArgument('--numUnits', {
      type: 'int',
      defaultValue: 32,
      help: 'Number of units in LSTM cell'
    });
    parser.addArgument('--returnSeq', {
      type: 'string',
      defaultValue: 'false',
      help: 'Whether to return sequences in the output of LSTM layer'
    });
    parser.addArgument(
        '--gpu', {action: 'storeTrue', help: 'Use GPU for training'});
    parser.addArgument('--optimizer', {
      type: 'string',
      defaultValue: 'adam',
      help: 'Optimizer to be used for model training'
    });
    parser.addArgument(
        '--epochs',
        {type: 'int', defaultValue: 10, help: 'Number of training epochs'});
    parser.addArgument(
        '--batchSize',
        {type: 'int', defaultValue: 128, help: 'Batch size for training'});
    parser.addArgument('--validationSplit', {
      type: 'float',
      defaultValue: 0.2,
      help: 'Validation split for training'
    });
    parser.addArgument('--modelSaveDir', {
      type: 'string',
      defaultValue: 'public/resources',
      help: 'Optional path for model saving.'
    });
    parser.addArgument('--embeddingFilesPrefix', {
      type: 'string',
      defaultValue: 'imdb/tmp/embed',
      help: 'Optional path prefix for saving embedding files that ' +
      'can be loaded in the Embedding Projector ' +
      '(https://projector.tensorflow.org/). For example, if this flag '  +
      'is configured to the value /tmp/embed, then the embedding vectors ' +
      'file will be written to /tmp/embed_vectors.tsv and the labels ' +
      'file will be written to /tmp/embed_label.tsv'
    });
    parser.addArgument('--logDir', {
      type: 'string',
      defaultValue:'imdb/tmp',
      help: 'Optional tensorboard log directory, to which the loss and ' +
      'accuracy will be logged during model training.'
    });
    parser.addArgument('--logUpdateFreq', {
      type: 'string',
      defaultValue: 'batch',
      optionStrings: ['batch', 'epoch'],
      help: 'Frequency at which the loss and accuracy will be logged to ' +
      'tensorboard.'
    });
    return parser.parseArgs();
  }
  
  async function main() {
    const args = parseArguments();
  
    let tfn;
    if (args.gpu) {
      console.log('Using GPU for training');
      tfn = require('@tensorflow/tfjs-node-gpu');
    } else {
      console.log('Using CPU for training');
      tfn = require('@tensorflow/tfjs-node');
    }
  
    console.log('Loading data...');
    const multihot = args.modelType === 'multihot';
    const {xTrain, yTrain, xTest, yTest} =
        await loadData(args.numWords, args.maxLen, multihot);
  
    console.log('Building model...');
    const model = buildModel(
        args.modelType, args.maxLen, args.numWords, args.embeddingSize,
         args.numUnits, args.returnSeq);
  
    model.compile({
      loss: 'binaryCrossentropy',
      optimizer: args.optimizer,
      metrics: ['acc']
    });
    model.summary();
  
    console.log('Training model...');
    await model.fit(xTrain, yTrain, {
      epochs: args.epochs,
      batchSize: args.batchSize,
      validationSplit: args.validationSplit,
      callbacks: args.logDir == null ? null : tfn.node.tensorBoard(args.logDir, {
        updateFreq: args.logUpdateFreq
      })
    });
  
    console.log('Evaluating model...');
    const [testLoss, testAcc] =
        model.evaluate(xTest, yTest, {batchSize: args.batchSize});
    console.log(`Evaluation loss: ${(await testLoss.data())[0].toFixed(4)}`);
    console.log(`Evaluation accuracy: ${(await testAcc.data())[0].toFixed(4)}`);
  
    // Save model.
    let metadata;
    if (args.modelSaveDir != null && args.modelSaveDir.length > 0) {
      if (multihot) {
        console.warn(
            'Skipping saving of multihot model, which is not supported.');
      } else {
        // Create base directory first.
        shelljs.mkdir('-p', args.modelSaveDir);
  
        // Load metadata template.
        console.log('Loading metadata template...');
        metadata = await loadMetadataTemplate();
  
        // Save metadata.
        metadata.epochs = args.epochs;
        metadata.embedding_size = args.embeddingSize;
        metadata.max_len = args.maxLen;
        metadata.model_type = args.modelType;
        metadata.batch_size = args.batchSize;
        metadata.vocabulary_size = args.numWords;
        const metadataPath = path.join(args.modelSaveDir, 'metadata.json');
        fs.writeFileSync(metadataPath, JSON.stringify(metadata));
        console.log(`Saved metadata to ${metadataPath}`);
  
        // Save model artifacts.
        await model.save(`file://${args.modelSaveDir}`);
        console.log(`Saved model to ${args.modelSaveDir}`);
      }
    }
  
    if (args.embeddingFilesPrefix != null &&
        args.embeddingFilesPrefix.length > 0) {
      if (metadata == null) {
        metadata = await loadMetadataTemplate();
      }
      await writeEmbeddingMatrixAndLabels(
          model, args.embeddingFilesPrefix, metadata.word_index,
          metadata.index_from);
    }
  }
  
  if (require.main === module) {
    main();
  }
  