Function extractEmbeddingMatrix

sentiment/embedding.js:34–46

* Extract the first embedding matrix from a TensorFlow.js model. * * @param {tf.model} model An instance of tf.Model, assumed to contain an * Embedding layer. * @retuns {tf.Tensor} The embedding matrix from the first Embedding * layer encoutnered while iterating through all layers of the mo

(model)

 * @throws Error if no embedding layer can be found in the model.
 */
function extractEmbeddingMatrix(model) {
  for (const layer of model.layers) {
    if (layer.getClassName() === 'Embedding') {
      const embed = layer.getWeights()[0];
      tf.util.assert(
        embed.rank === 2,
        `Expected the rank of an embedding matrix to be 2, ` + 
        `but got ${embed.rank}`);
      return embed;
    }
  }
  throw new Error('Cannot find any Embedding layer in model.');
}

/**
 * Write the values of the first embedding matrix of a model to files.