Function buildModel

sentiment/train.js:36–85

(modelType, maxLen, vocabularySize, embeddingSize)

 * @returns An uncompiled instance of `tf.Model`.
 */
export function buildModel(modelType, maxLen, vocabularySize, embeddingSize) {
  // 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: 32}));
    } else if (modelType === 'lstm') {
      model.add(tf.layers.lstm({units: 32}));
    } 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}`);
    }
  }
  model.add(tf.layers.dense({units: 1, activation: 'sigmoid'}));
  return model;
}

function parseArguments() {
  const parser = new ArgumentParser(