Class LSTMTextGenerator

lstm-text-generation/index.js:44–129

 * - Generating random text using the LSTM model.
 */
export class LSTMTextGenerator {
  /**
   * Constructor of NeuralNetworkTextGenerator.
   *
   * @param {TextData} textData An instance of `TextData`.
   */
  constructor(textData) {
    this.textData_ = textData;
    this.charSetSize_ = textData.charSetSize();
    this.sampleLen_ = textData.sampleLen();
    this.textLen_ = textData.textLen();
  }

  /**
   * Create LSTM model from scratch.
   *
   * @param {number | number[]} lstmLayerSizes Sizes of the LSTM layers, as a
   *   number or an non-empty array of numbers.
   */
  createModel(lstmLayerSizes) {
    this.model = model.createModel(
        this.sampleLen_, this.charSetSize_, lstmLayerSizes);
  }

  /**
   * Compile model for training.
   *
   * @param {number} learningRate The learning rate to use during training.
   */
  compileModel(learningRate) {
    model.compileModel(this.model, learningRate);
  }

  /**
   * Train the LSTM model.
   *
   * @param {number} numEpochs Number of epochs to train the model for.
   * @param {number} examplesPerEpoch Number of epochs to use in each training
   *   epochs.
   * @param {number} batchSize Batch size to use during training.
   * @param {number} validationSplit Validation split to be used during the
   *   training epochs.
   */
  async fitModel(numEpochs, examplesPerEpoch, batchSize, validationSplit) {
    let batchCount = 0;
    const batchesPerEpoch = examplesPerEpoch / batchSize;
    const totalBatches = numEpochs * batchesPerEpoch;
    let t = new Date().getTime();

    onTrainBegin();
    const callbacks = {
      onBatchEnd: async (batch, logs) => {
        // Calculate the training speed in the current batch, in # of
        // examples per second.
        const t1 = new Date().getTime();
        const examplesPerSec = batchSize / ((t1 - t) / 1e3);
        t = t1;
        onTrainBatchEnd(logs, ++batchCount / totalBatches, examplesPerSec);