用NodeJS/TensorFlowJS调用BERT实现文本分类

$ curl --output train.json https://deepdialog.coding.net/p/dataset/d/dataset/git/raw/master/LCQMC/train.json
$ curl --output dev.json https://deepdialog.coding.net/p/dataset/d/dataset/git/raw/master/LCQMC/dev.json

$ curl --output vocab.txt https://deepdialog.coding.net/p/zh-roberta-wwm/d/zh-roberta-wwm/git/raw/master/vocab.txt

$ mkdir -p bert
$ cd bert
$ curl --output bert.tar.gz https://deepdialog.coding.net/p/zh-roberta-wwm/d/zh-roberta-wwm/git/raw/master/zh-roberta-wwm-L12.tar.gz
$ tar xvzf bert.tar.gz
$ cd ..

npm i install @tensorflow/tfjs-node tokenizers

const fs = require('fs')
const tf = require('@tensorflow/tfjs-node')
// huggingface的bert分词包
const BertWordPieceTokenizer = require('tokenizers').BertWordPieceTokenizer


/**
 * 构建文本分类模型
 * 输入的是BERT输出的sequence_output序列
 * 输出2分类softmax
 */
function buildModel() {
    const input = tf.input({shape: [null, 768], dtype: 'float32'})
    // 这里之所以用rnn对bert输出序列进行训练，而不是直接针对[CLS]输出进行训练
    // 是因为如果不fine-tune bert的参数的话，只用[CLS]效果会差一点
    const rnn = tf.layers.bidirectional({
        layer: tf.layers.lstm({units: 128, returnSequences: false})
    })
    // masking很重要，我封装的模型padding的部分会输出 0.0 （有可能是 -0.0，但是也可以被mask）
    const mask = tf.layers.masking({maskValue: 0.0})
    const dense = tf.layers.dense({units: 2, activation: 'softmax'})
    const output = dense.apply(rnn.apply(mask.apply(input)))
    const model = tf.model({inputs: input, outputs: output})
    model.compile({
        optimizer: 'adam',
        loss: 'sparseCategoricalCrossentropy',
        metrics: ['acc'],
    })
    return model
}


(async () => {

    // 加载词表/分词和BERT
    const wordPieceTokenizer = await BertWordPieceTokenizer.fromOptions({ vocabFile: "./vocab.txt" })
    const bert = await tf.node.loadSavedModel('./bert')

    // 构建数据流
    // 文本输入会经过tokenizers
    // 然后用bert计算出sequence_output
    // 不更新bert的参数是因为nodejs现在还无法训练读取的模型
    function makeGenerator(objs, batchSize) {
        function* dataGenerator() {
            let xs = []
            let ys = []
            for (const obj of objs) {
                xs.push(obj['tokens'])
                ys.push(Number.parseInt(obj['label']))
                if (xs.length == ys.length && xs.length == batchSize) {
                    // 下面几行，是对数据进行padding到一样长度，补足的部分使用空字符串
                    const maxLength = Math.max.apply(
                        Math,
                        xs.map(x => x.length)
                    )
                    xs = xs.map(x => {
                        while(x.length < maxLength) {
                            x = x.concat([''])
                        }
                        return x
                    })
                    xs = tf.tensor(xs)
                    // 这一步是得到bert的输出结果
                    // 如果输入是dict格式，输出也会是dict格式，可以参考tfjs的源代码
                    // 这一步也可以单独用，就类似bert-as-a-service一样
                    xs = bert.predict({
                        input_1: xs
                    })['sequence_output']
                    ys = tf.tensor(ys)
                    // bert的输出作为文本分类模型的输入(xs)
                    // 标签作为文本分类模型的目标(ys)
                    yield {xs, ys}
                    xs = []
                    ys = []
                }
            }
        }
        return dataGenerator
    }

    // 数据集，格式是jsonl，所以用这种方法读取
    console.log('Read dataset')
    const trainObjs = fs.readFileSync(
        'train.json',
        {encoding: 'utf-8'}
    ).split(/\n/).map(JSON.parse)
    const devObjs = fs.readFileSync(
        'dev.json',
        {encoding: 'utf-8'}
    ).split(/\n/).map(JSON.parse)

    // 这里先对分词，是因为分词是async异步函数，而tensorflowjs的generator不支持异步yield
    console.log('Tokenize train dataset')
    for (const obj of trainObjs) {
        obj['tokens'] = (await wordPieceTokenizer.encode(
            obj['sentence1'], obj['sentence2']
        )).tokens
    }
    console.log('Tokenize dev dataset')
    for (const obj of devObjs) {
        obj['tokens'] = (await wordPieceTokenizer.encode(
            obj['sentence1'], obj['sentence2']
        )).tokens
    }
    console.log('Start training')

    const batchSize = 32
    const dsTrain = tf.data.generator(makeGenerator(trainObjs, batchSize)).repeat()
    const dsDev = tf.data.generator(makeGenerator(devObjs, batchSize)).repeat()
    const model = buildModel()
    model.fitDataset(dsTrain, {
        batchesPerEpoch: Math.floor(trainObjs.length / batchSize),
        epochs: 1,
        batch_size: batchSize,
        validationData: dsDev,
        validationBatches: Math.floor(devObjs.length / batchSize),
    })

    model.evaluateDataset(dsDev, {
        batches: Math.floor(devObjs.length / batchSize),
    })

})()