Search code examples
javascripttensorflowweb-workertensorflow.js

TensorflowJS universal sentence encoder in WebWorker

Trying to run USE in webworker to avoid performance issues, but it lacks embed() method.

WebWorker

importScripts('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.2.9/dist/tf.min.js');
importScripts('https://cdn.jsdelivr.net/npm/@tensorflow-models/universal-sentence-encoder@1.1.1/dist/universal-sentence-encoder.min.js');
tf.setBackend('cpu');

onmessage = message => {
  use.load().then(model => {
    const sentences = [
      'I like my phone.',
      'Your cellphone looks great.'
    ];
    /**
      model = { model, tokenizer }
    */
    model.embed(sentences).then(async embeddings => {
      const vec = await embeddings.array();
      const cosine = tf.losses.cosineDistance(vec[0], vec[1], 0);
      const result = await cosine.data();
      console.log(`${Math.round((1 - result) * 100)}%`);
    });
  });
}

There is no embed() inside model which is the case for main thread
Is it even possible to run USE in webworkers?

Solution

  • Universal Sentence Encoder (USE) can be used in a webworker. See below this working snippet with the result sent from the worker to the main thread

    <head>
	<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@0.14.2/dist/tf.min.js"></script>
	<script>
        const worker_function = () =>
{
	importScripts('https://cdn.jsdelivr.net/npm/@tensorflow/tfjs@1.2.9/dist/tf.min.js');
	importScripts('https://cdn.jsdelivr.net/npm/@tensorflow-models/universal-sentence-encoder@1.1.1/dist/universal-sentence-encoder.min.js');
	tf.setBackend('cpu');
	onmessage = message =>
	{
        console.log("from worker")
		use.load().then(model =>
		{
			const sentences = [
				'I like my phone.',
				'Your cellphone looks great.'
			];
			model.embed(sentences).then(async embeddings =>
			{ 
                console.log(await embeddings)
				const vec = await embeddings.array();
				const cosine = tf.losses.cosineDistance(vec[0], vec[1], 0);
				const result = await cosine.data();
                console.log(`${Math.round((1 - result) * 100)}%`);
                
                postMessage({res: cosine.dataSync(), shape: cosine.shape})
            });
            
		});
	};
}
if (window != self)
	worker_function();
    </script>
	<script>
        const worker = new Worker(URL.createObjectURL(new Blob(["(" + worker_function.toString() + ")()"], { type: 'text/javascript' })));
        worker.postMessage({});
        worker.onmessage = (message) => {
            console.log('from main thread')
            const {data} = message
        	tf.tensor(message.data.res, message.data.shape).print()
        }
    </script>
</head>