NLP Keras model in browser with TensorFlow.js
How to use your Keras model in browser with tensorflow.js
In this article I’ll attempt to cover three things:
- How to write simple Named-Entity Recognition model — typical Natural Language Processing (NLP) task.
- How to export this model to TensorFlow.js format.
- How to make simple WEB application for searching Named-Entity in string without back-end.
You can test it on my personal site.
Firstly, this blog post for people who know the main stages of deep learning (RNN, Attention) and who know basic of JavaScript. If you are not good with RNN, I recommend you tread “The Unreasonable Effectiveness of Recurrent Neural Networks” by Andrej Karpathy.
TensorFlow.js is a JavaScript library for developing and training ML models in JavaScript, and deploying in a browser or on Node.js.
For this example, we will use simple keras model for solving the classic NER task. We will train on CoNLL2003 data set. Our model it’s just word embedding, GRU and very simple attention mechanism. After we will visualize the attention vector in the browser. If you are familiar with modern approaches for solving a similar task, you know, that this approach is not state of the art. However, for running it in a browser and just, for example, it enough. Furthermore, if you are familiar with NER task, you can skip the next two sections about Named-Entity recognition, data and NER neural network model on Keras and look at the source code.
Task and Data
Depending on your experience, you may have knowledge of it under different names, like a sequence tagging, Part-of-Speech tagging or, like in our task — Named-Entity recognition.
In generally, Named-Entity Recognition (NER) task is seq2seq task. For each token from x_i we have to predict token y_i, where x is an input sequence and y is a sequence of named entities.
In this example, we will be looking for persons (B-PER, I-PER), locations (B-LOC, I-LOC) and organizations (B-ORG, I-ORG), in addition, the model will define special entity — MISC — named entity but not a person, location and organization.
Firstly, we have to prepare data for a computer (yes, we will use a computer for solving this task :) ).
In this blog post, we don’t have a goal to make SOTA result on CoNLL2003 data set, that’s why our pre-processing isn’t really good. But just an example, we will load data by load_data method:
As you know, neural networks can’t work with words, only with numbers. That’s why we should represent words as numbers. It’s not a difficult task, we can enumerate all unique words and write the number of word instead word. For storing numbers and words, we can make vocabulary. This vocabulary should support words “unknown” (<UNK>) because when we will make a prediction for a new string with new words that are not in the vocabulary. Word “padded” (<PAD>) because for the neural network, all strings should have the same size, and when some string will smaller than another, we fill the gaps with this word.
In addition, let’s write the simple helper for translate sentence to the sequence of numbers.
As you can read above, we should pad sequences for working with neural network, to do this, you can use internal Keras method pad_sequences.
Model
Let me guess…RNN ?
Yes, it’s RNN. More specifically, it’s GRUs (Gated Recurrent Units) and simple attention layer. As a word representation used GloVe. In this blog post, I will not go into the details and just leave the code of model here. I hope it’s easy to understand.
After building the model, we should compile, train and save the model. You can guess, for running this model in a browser, we should save not only weights for the model but also model description and vocabularies for words and tags. Let’s define the method for the export model and vocabularies to JavaScript supported format (JSON basically).
Finally, let’s compile, train and export model:
Full code for this steps, you can find in my GitHub repository in train.py.
Development environment
Well, the model is ready and now, we should start developing a web application to evaluating this mode in a browser. We need a development environment. Basically, it’s not important, how you will store your model, weights, and vocabulary, but for “easy start”, I will show you, my solution — node.js local server.
We need two files: package.json and server.js.
in server.js we defined the static folder for storing model, js scripts and all other files. For using this server, you should enter
npm install && node server.js
in your terminal. After that, you can get access to your files in the browser by http://localhost:8081
Web application
There are index.html, predict.js and files that were generated in previous step. As you see it’s a very small web application. index.html contains requirements and the input field for entering strings.
Now, the most interesting part of the blog post — about tensorflow.js. You can load model by tf.loadLayersModel method calling await operator. It’s important because we don’t want to block our web application while loading model. If we load model, we will get the model that can predict tokens only, but what about attention vector? How to get activations from attention? For getting data from internal layers, in tensorflow.js we should make a new model which outputs layers will contain outputs and other layers from the original model, like this:
Here, model is the original model, emodel is model with the attention_vector as outputs.
Preprocessing
Now we have to implement preprocessing for strings as we did in our Python script. It’s a not a difficult task for us because the regular expression on Python and JavaScript is very similar and other method are similar too.
Make predictions
Now we should provide transfer data from simple text string format to TF format — tensor. In the previous section, we wrote helpers to transfer string to an array of numbers. Now we should create tf.tensor from this array. As you remember, the input layer of the model has the shape (None, 113), that’s why we should expand the dimension of input tensor. Well, it’s all, now we can make a prediction in the browser by .predict method. After that, you should print predicted data into a browser and your web application with neural network without back-end is ready.
Conclusion
TensorFlow.js is a library for using neural networks in a browser like a chrome, Firefox or safari. If you open this web app in your iPhone or Android smartphone, it will work too.
You can test it on my personal web site by this link: http://deepdivision.net/2019/03/20/nlp-keras-model-in-browser-with-tensorflowjs.html
In addition, you can find this code with some additions in my GitHub.
