Public
Edited
Apr 29, 2024
Insert cell
Insert cell
tf = require('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs/dist/tf.min.js')
Insert cell
results = {
// Define a function that creates the Word2Vec model and trains it
async function trainWord2Vec() {
let output = {};

// Define some sample sentences
const sentences = [
['I', 'love', 'machine', 'learning'],
['Deep', 'learning', 'is', 'amazing'],
['Artificial', 'intelligence', 'is', 'the', 'future']
];

output.sentences = sentences;

// Create word2vec model
const model = tf.sequential();
model.add(tf.layers.dense({ inputShape: [sentences[0].length], units: 2 }));
model.compile({ optimizer: 'sgd', loss: 'meanSquaredError' });

// Convert words to indices
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;

// Convert sentences to input data
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 };

// Convert input data to tensors
const xsTensor = tf.tensor2d(xs);
const ysTensor = tf.tensor2d(ys);

output.tensors = { xsTensor, ysTensor };

// Train the model
// await model.fit(xsTensor, ysTensor, { epochs: 100 });

// Test the model
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;
}
}
//const testOutput = model.predict(tf.tensor2d([testInput]));
// output.results = testOutput.dataSync();

return output;
}

// Call the function to train the model and get the results
const results = await trainWord2Vec();

return results
}
Insert cell

Purpose-built for displays of data

Observable is your go-to platform for exploring data and creating expressive data visualizations. Use reactive JavaScript notebooks for prototyping and a collaborative canvas for visual data exploration and dashboard creation.
Learn more