ml5 Example: Word2Vec

// Copyright (c) 2018 ml5

//

// This software is released under the MIT License.

// https://opensource.org/licenses/MIT

​

/* ===

ml5 Example

Word2Vec example with p5.js. Using a pre-trained model on common English words.

=== */

​

// Create the Word2Vec model with pre-trained file of 10,000 words

const word2Vec = ml5.word2vec('https://cdn.rawgit.com/ml5js/ml5-examples/master/p5js/Word2Vec/data/wordvecs10000.json', modelLoaded);

let modelReady = false;

​

function setup() {

  noLoop();

  //noCanvas();

  createCanvas(windowWidth, windowHeight);

  // Select all the DOM elements

  let nearWordInput = createInput('rainbow');

  nearWordInput.position(width/6, height/3);

  let nearButton = createButton('is nearest to ...');

  nearButton.position(width/6, height/3+50);

  let nearResults = createP('');

  nearResults.position(width/6, height/3+100);

  text('Between', width*2/6, height/3);

  let betweenWordInput1 = createInput('rainbow');

  betweenWordInput1.position(width*2/6, height/3+50);

  text('and', width*2/6, height/3+100);

  let betweenWordInput2 = createInput('kitten');

  betweenWordInput2.position(width*2/6, height/3+150);

  let betweenButton = createButton("is ...");

  betweenButton.position(width*2/6, height/3+200);

  let betweenResults = createP("");

  betweenResults.position(width*2/6, height/3+250);

​

  let addInput1 = createInput("pizza");

  addInput1.position(width*3/6, height/3);

  text('is to', width*3/6, height/3+50);

  let addInput2 = createInput("humans");

  addInput2.position(width*3/6, height/3+100);

  text('as', width*3/6, height/3+150);

  let addInput3 = createInput("fish");

  addInput3.position(width*3/6, height/3+200);

  let addButton = createButton("is to ...");

  addButton.position(width*3/6, height/3+250);

  let addResults = createP("");

  addResults.position(width*3/6, height/3+300);

​

  // Finding the nearest words

  nearButton.mousePressed(() => {

    let word = nearWordInput.value();

    nearResults.html(findNearest(word, 10));

  });

​

  // Findind the average of two words

  betweenButton.mousePressed(() => {

    let word1 = betweenWordInput1.value();

    let word2 = betweenWordInput2.value();

    let average = word2Vec.average([word1, word2], 1);

    betweenResults.html(average[0].vector);

  });

​

  // Adding two words together to "solve" an analogy

  addButton.mousePressed(() => {

    let is1 = addInput1.value();

    let to1 = addInput2.value();

    let is2 = addInput3.value();

    let difference = word2Vec.subtract([to1, is1]);

    let to2 = word2Vec.add([is2, difference[0].vector]);

    addResults.html(to2[0].vector);

  });

}

​

// Model is ready

function modelLoaded() {

  print("model loaded");

  modelReady = true;

}

​

// Converts the result of nearest() to html to display

function findNearest(word, n = 10) {

  if (modelReady) {

    let nearest = word2Vec.nearest(word, n);

    if (!nearest) {

      return 'No word vector found';

    }

    let output = '';

    for (let i = 0; i < nearest.length; i++) {

      output += nearest[i].vector + '<br/>';

    }

    return output;

  }

  return 'Model has not loaded yet!';

}