openvino/samples/js/node/notebooks/tokens_bert.js

117 lines
2.6 KiB
JavaScript

const fs = require('node:fs/promises');
exports.cleanWord = cleanWord;
exports.encodeByVoc = encodeByVoc;
exports.textToTokens = textToTokens;
exports.splitToWords = splitToWords;
exports.loadVocabFile = loadVocabFile;
// Load vocabulary file for encoding
async function loadVocabFile(vocabFileName) {
const vocab = {};
const lines = (await fs.readFile(vocabFileName, 'utf-8')).split('\n');
lines.forEach((line, index) => {
const token = line.trim();
vocab[token] = index;
});
return vocab;
}
// Remove mark and control chars
function cleanWord(w) {
let wo = ''; // Accumulator for output word
for (const c of w.normalize('NFD')) {
const charCode = c.charCodeAt(0);
// Remove mark nonspacing code and controls
if (charCode < 32 || charCode == 127) continue;
wo += c;
}
return wo;
}
// Split word by vocab items and get tok codes iteratively return codes
function encodeByVoc(w, vocab) {
w = cleanWord(w);
const res = [];
const wordIndexes = splitToWords(w);
for (let el of wordIndexes) {
const [s0, e0] = el;
let s = s0;
let e = e0;
const tokens = [];
while (e > s) {
const subword = s == s0 ? w.slice(s, e) : '##' + w.slice(s, e);
if (vocab[subword]) {
tokens.push(vocab[subword]);
s = e;
e = e0;
}
else e -= 1;
}
if (s < e0) tokens.push(vocab['[UNK]']);
res.push(...tokens);
}
return res;
}
// Split big text into words by spaces
// Return start and end indexes of words
function splitToWords(text) {
let start;
let prevIsSep = true; // Mark initial prev as space to start word from 0 char
const result = [];
for (let i = 0; i < text.length + 1; i++) {
const c = text[i] || ' ';
const isPunc = /[!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~]/.test(c);
const curIsSep = c.trim() === '' || isPunc;
if (prevIsSep !== curIsSep)
if (prevIsSep)
start = i;
else {
result.push([start, i]);
prevIsSep = curIsSep;
}
if (isPunc) result.push([i, i + 1]);
prevIsSep = curIsSep;
}
return result;
}
// Get big text and return list of token id and start-end positions
// for each id in original texts
function textToTokens(text, vocab) {
const tokensId = [];
const tokensSe = [];
const wordIndices = splitToWords(text);
for (const [start, end] of wordIndices) {
const word = text.slice(start, end);
const encodedTokens = encodeByVoc(word, vocab);
for (const token of encodedTokens) {
tokensId.push(token);
tokensSe.push([start, end]);
}
}
return [tokensId, tokensSe];
}