Generate text with character-based RNN

tensorflow.org has a nice tutorial on generating text using a character-based RNN. I wanted to implement the architecture in Javascript and make it interactive, but it turns out training a deep neural network is very computationally intensive, and I cannot make it reasonably interactive with simple Javascript code...

Nevertheless, the following is a toy that lets you interactively train a RNN model and make it generate text.

Pressing the 'Train ONE epoch' button trains the model with one pass of the Training Corpus. More than 10 epochs are needed to make the model generate something other than what looks like a wall of random letters.

const vocabSize = 61;
const embedSize = 16;
const numHiddens = 16;

forwardWithState(x: nd.T, h: nd.T, training: boolean): [nd.T, nd.T] {
  const [batchSize, seqLength] = x.shape();

  // Embedding layer
  x = nd.reshape(x, [batchSize * seqLength]);
  x = this.layerEmbedding.forward(x, training) as nd.T;

  // RNN layer
  x = nd.reshape(x, [batchSize, seqLength, this.embedSize]);
  x = nd.transpose(x, [1, 0, 2]);
  let nextH: nd.T;
  [x, nextH] = this.layerRNN.forward([x, h], training) as [nd.T, nd.T];
  x = nd.transpose(x, [1, 0, 2]);

  // FC outout layer
  x = this.layerDense.forward(x, training) as nd.T;
  x = nd.reshape(x, [batchSize, seqLength, this.vocabSize]);

  return [x, nextH];
}