Add tests for Word2Vec

- validates that there are no leaked tensors
- validates basic functionality of nearest
- validates that add, subtract and average return things

- Fixes memory leaks in add, subtract, average, and addOrSubtract functions
- Adds a general dispose to the Word2Vec class

src/Word2vec/index.js

@@ -34,20 +34,31 @@ class Word2Vec {
  });
  }
 
+ dispose() {
+ Object.values(this.model).forEach(x => x.dispose());
+ }
+
  add(inputs, max = 1) {
- const sum = Word2Vec.addOrSubtract(this.model, inputs, 'ADD');
- return Word2Vec.nearest(this.model, sum, inputs.length, inputs.length + max);
+ return tf.tidy(() => {
+ const sum = Word2Vec.addOrSubtract(this.model, inputs, 'ADD');
+ console.log(sum);
+ return Word2Vec.nearest(this.model, sum, inputs.length, inputs.length + max);
+ });
  }
 
  subtract(inputs, max = 1) {
- const subtraction = Word2Vec.addOrSubtract(this.model, inputs, 'SUBTRACT');
- return Word2Vec.nearest(this.model, subtraction, inputs.length, inputs.length + max);
+ return tf.tidy(() => {
+ const subtraction = Word2Vec.addOrSubtract(this.model, inputs, 'SUBTRACT');
+ return Word2Vec.nearest(this.model, subtraction, inputs.length, inputs.length + max);
+ });
  }
 
  average(inputs, max = 1) {
- const sum = Word2Vec.addOrSubtract(this.model, inputs, 'ADD');
- const avg = tf.div(sum, tf.tensor(inputs.length));
- return Word2Vec.nearest(this.model, avg, inputs.length, inputs.length + max);
+ return tf.tidy(() => {
+ const sum = Word2Vec.addOrSubtract(this.model, inputs, 'ADD');
+ const avg = tf.div(sum, tf.tensor(inputs.length));
+ return Word2Vec.nearest(this.model, avg, inputs.length, inputs.length + max);
+ });
  }
 
  nearest(input, max = 10) {
@@ -64,34 +75,36 @@ class Word2Vec {
  }
 
  static addOrSubtract(model, values, operation) {
- const vectors = [];
- const notFound = [];
- if (values.length < 2) {
- throw new Error('Invalid input, must be passed more than 1 value');
- }
- values.forEach((value) => {
- const vector = model[value];
- if (!vector) {
- notFound.push(value);
- } else {
- vectors.push(vector);
+ return tf.tidy(() => {
+ const vectors = [];
+ const notFound = [];
+ if (values.length < 2) {
+ throw new Error('Invalid input, must be passed more than 1 value');
  }
- });
+ values.forEach((value) => {
+ const vector = model[value];
+ if (!vector) {
+ notFound.push(value);
+ } else {
+ vectors.push(vector);
+ }
+ });
 
- if (notFound.length > 0) {
- throw new Error(`Invalid input, vector not found for: ${notFound.toString()}`);
- }
- let result = vectors[0];
- if (operation === 'ADD') {
- for (let i = 1; i < vectors.length; i += 1) {
- result = tf.add(result, vectors[i]);
+ if (notFound.length > 0) {
+ throw new Error(`Invalid input, vector not found for: ${notFound.toString()}`);
  }
- } else {
- for (let i = 1; i < vectors.length; i += 1) {
- result = tf.sub(result, vectors[i]);
+ let result = vectors[0];
+ if (operation === 'ADD') {
+ for (let i = 1; i < vectors.length; i += 1) {
+ result = tf.add(result, vectors[i]);
+ }
+ } else {
+ for (let i = 1; i < vectors.length; i += 1) {
+ result = tf.sub(result, vectors[i]);
+ }
  }
- }
- return result;
+  return result;
+ });
  }
 
  static nearest(model, input, start, max) {

src/Word2vec/index_test.js

@@ -1,19 +1,96 @@
-// import Word2Vec from './index';
-
-// const URL = 'https://raw.githubusercontent.com/ml5js/ml5-examples/master/p5js/07_Word2Vec/data/wordvecs1000.json';
-
-// describe('initialize word2vec', () => {
-// let word2vec;
-// beforeAll((done) => {
-// // word2vec = new Word2Vec(URL);
-// done();
-// });
-
-// // it('creates a new instance', (done) => {
-// // expect(word2vec).toEqual(jasmine.objectContaining({
-// // ready: true,
-// // modelSize: 1,
-// // }));
-// // done();
-// // });
-// });
+const { tf, word2vec } = ml5;
+
+const URL = 'https://raw.githubusercontent.com/ml5js/ml5-examples/master/p5js/Word2Vec/data/wordvecs1000.json';
+
+describe('word2vec', () => {
+ let word2vecInstance;
+ let numTensorsBeforeAll;
+ let numTensorsBeforeEach;
+ beforeAll((done) => {
+ numTensorsBeforeAll = tf.memory().numTensors;
+ word2vecInstance = word2vec(URL, done);
+ });
+
+ afterAll(() => {
+ word2vecInstance.dispose();
+ let numTensorsAfterAll = tf.memory().numTensors;
+ if(numTensorsBeforeAll !== numTensorsAfterAll) {
+ throw new Error(`Leaking Tensors (${numTensorsAfterAll} vs ${numTensorsBeforeAll})`);
+ }
+ });
+
+ beforeEach(() => {
+ numTensorsBeforeEach = tf.memory().numTensors;
+ });
+
+ afterEach(() => {
+ let numTensorsAfterEach = tf.memory().numTensors;
+ if(numTensorsBeforeEach !== numTensorsAfterEach) {
+ throw new Error(`Leaking Tensors (${numTensorsAfterEach} vs ${numTensorsBeforeEach})`);
+ }
+ });
+
+ it('creates a new instance', () => {
+ expect(word2vecInstance).toEqual(jasmine.objectContaining({
+ ready: true,
+ modelSize: 1,
+ }));
+ });
+
+ describe('getRandomWord', () => {
+ it('returns a word', () => {
+ let word = word2vecInstance.getRandomWord();
+ expect(typeof word).toEqual('string');
+ });
+ });
+
+ describe('nearest', () => {
+ it('returns a sorted array of nearest words', () => {
+ for(let i = 0; i < 100; i++) {
+ let word = word2vecInstance.getRandomWord();
+ let nearest = word2vecInstance.nearest(word);
+ let currentDistance = 0;
+ for(let { word, distance: nextDistance } of nearest) {
+ expect(typeof word).toEqual('string');
+ expect(nextDistance).toBeGreaterThan(currentDistance);
+ currentDistance = nextDistance;
+ }
+ }
+ });
+
+ it('returns a list of the right length', () => {
+ for(let i = 0; i < 100; i++) {
+ let word = word2vecInstance.getRandomWord();
+ let nearest = word2vecInstance.nearest(word, i);
+ expect(nearest.length).toEqual(i);
+ }
+ });
+ });
+
+ describe('add', () => {
+ it('returns a value', () => {
+ let word1 = word2vecInstance.getRandomWord();
+ let word2 = word2vecInstance.getRandomWord();
+ let sum = word2vecInstance.subtract([word1, word2]);
+ expect(sum[0].distance).toBeGreaterThan(0);
+ })
+ });
+
+ describe('subtract', () => {
+ it('returns a value', () => {
+ let word1 = word2vecInstance.getRandomWord();
+ let word2 = word2vecInstance.getRandomWord();
+ let sum = word2vecInstance.subtract([word1, word2]);
+ expect(sum[0].distance).toBeGreaterThan(0);
+ })
+ });
+
+ describe('average', () => {
+ it('returns a value', () => {
+ let word1 = word2vecInstance.getRandomWord();
+ let word2 = word2vecInstance.getRandomWord();
+ let average = word2vecInstance.average([word1, word2]);
+ expect(average[0].distance).toBeGreaterThan(0);
+ });
+ });
+});