tiny-tensorflow/predict.js

async function run() {
    const { vitchen_toxicity_encode } = require('./toxicity-tfjs/bert-tokenizer');
    const {loadSavedModel, tensor} = require('./toxicity-tfjs/node');
    const model = await loadSavedModel('./bert-small-multilingual');
    const input = vitchen_toxicity_encode(`What's up?`);
    let t = tensor(input, [192], 'int32');
    const prediction = model.predict({
        input_ids: t
    })['output_0'];
    console.log(prediction.dataSync()[0])
}
run()