results = {
async function trainWord2Vec() {
let output = {};
const sentences = [
['I', 'love', 'machine', 'learning'],
['Deep', 'learning', 'is', 'amazing'],
['Artificial', 'intelligence', 'is', 'the', 'future']
];
output.sentences = sentences;
const model = tf.sequential();
model.add(tf.layers.dense({ inputShape: [sentences[0].length], units: 2 }));
model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' });
const word2idx = {};
let idx = 0;
for (const sentence of sentences) {
for (const word of sentence) {
if (!(word in word2idx)) {
word2idx[word] = idx++;
}
}
}
output.idx = word2idx;
output.model = model;
const xs = [];
const ys = [];
for (const sentence of sentences) {
const x = new Array(Object.keys(word2idx).length).fill(0);
const y = new Array(2).fill(0);
for (const word of sentence) {
x[word2idx[word]] = 1;
}
xs.push(x);
ys.push(y);
}
output.xy = { xs, ys };
const xsTensor = tf.tensor2d(xs);
const ysTensor = tf.tensor2d(ys);
output.tensors = { xsTensor, ysTensor };
const testSentence = ['Deep', 'learning', 'is', 'awesome'];
const testInput = new Array(Object.keys(word2idx).length).fill(0);
for (const word of testSentence) {
if (word in word2idx) {
testInput[word2idx[word]] = 1;
}
}
return output;
}
const results = await trainWord2Vec();
return results
}