Function main

lstm-text-generation/gen_node.js:88–124

()

}

async function main() {
  const args = parseArgs();

  if (args.gpu) {
    console.log('Using GPU');
    require('@tensorflow/tfjs-node-gpu');
  } else {
    console.log('Using CPU');
    require('@tensorflow/tfjs-node');
  }

  // Load the model.
  const loadModel = tf.loadModel || tf.loadLayersModel;
  const model = await loadModel(`file://${args.modelJSONPath}`);

  const sampleLen = model.inputs[0].shape[1];

  // Create the text data object.
  let localTextDataPath = args.textDatasetPath;
  if (args.textDatasetName) {
    const textDataURL = TEXT_DATA_URLS[args.textDatasetName].url;
    localTextDataPath = path.join(os.tmpdir(), path.basename(textDataURL));
    await maybeDownload(textDataURL, localTextDataPath);
  }
  const text = fs.readFileSync(localTextDataPath, {encoding: 'utf-8'});
  const textData = new TextData('text-data', text, sampleLen, args.sampleStep);

  // Get a seed text from the text data object.
  const [seed, seedIndices] = textData.getRandomSlice();
  
  console.log(`Seed text:\n"${seed}"\n`);

  const generated = await generateText(
      model, textData, seedIndices, args.genLength, args.temperature);

  console.log(`Generated text:\n"${generated}"\n`);
}

main();