Bringing State-of-the-Art Text Classification to Mobile Devices with TensorFlow Lite

ByJordan Brown

Text classification is a core natural language processing (NLP) task that has seen tremendous progress in recent years thanks to advances in deep learning. The ability to automatically categorize text has numerous applications, from sentiment analysis and content moderation to document routing and spam detection.

Traditionally, deploying NLP models has been limited to the cloud due to their computational cost. However, with the exponential growth of smartphones – over 6 billion worldwide as of 2022 – there is an increasing need for AI capabilities directly on edge devices. Benefits include:

  • Faster inference latency
  • Increased privacy and security
  • Reduced connectivity requirements
  • Personalization to each user‘s data

TensorFlow Lite is a powerful enabler for mobile NLP deployments. Released in 2017, it‘s a lightweight solution for running machine learning models on mobile, embedded and IoT devices. Models can be compressed up to 4x via techniques like quantization and pruning while maintaining accuracy. TensorFlow Lite has seen rapid adoption – it now has over 6 billion installs across 100,000+ apps.

In this tutorial, we‘ll walk through how to build a state-of-the-art text classification app on Android with TensorFlow Lite. We‘ll cover the full workflow from data collection to model training to deployment. Let‘s dive in!

Step 1: Collecting and Preparing Training Data

The first step in any machine learning project is gathering a high-quality, labeled dataset. For text classification, that means a corpus of text samples annotated with their target categories. Some example use cases:

  • News articles labeled by topic (e.g. sports, politics, technology)
  • Social media posts classified by sentiment (positive, negative, neutral)
  • Email messages tagged as spam or not spam
  • Product reviews categorized by department

The amount of data needed depends on the complexity of the task and model, but a good rule of thumb is to have at least 1,000 samples per class. Some potential sources:

Key considerations include:

  • Ensuring data diversity and coverage of important categories
  • Balancing class sizes to avoid bias
  • Establishing annotation guidelines for consistency
  • Performing quality control on labels
  • Preprocessing text to remove noise (e.g. HTML tags, emojis)
  • Splitting data into train, validation and test sets

It‘s important to allocate sufficient time for data collection and cleaning, as the quality of the dataset is a major factor in downstream model performance.

Step 2: Choosing a Model Architecture

With our training data ready, the next decision is selecting a model architecture. TensorFlow Lite supports several state-of-the-art networks that are well-suited for on-device text classification:

  • MobileBERT: A compact variant of the popular BERT model that achieves comparable performance with a 4.3x reduction in size and 5.5x latency speedup. Pre-trained on Wikipedia and the Toronto Book Corpus, it generates contextual word embeddings that can be fine-tuned for downstream tasks.

  • Word2Vec + CNN: A combination of Word2Vec embeddings and a convolutional neural network. Word2Vec captures semantic relationships between words, while the CNN learns to extract higher-level features. Faster but less accurate than BERT-based models.

  • FastText + BiLSTM: FastText embeddings paired with a bidirectional long short-term memory network. FastText extends Word2Vec with subword information to better handle out-of-vocabulary words. The BiLSTM encodes word order dependencies in both forward and backward directions.

The model choice depends on dataset size, inference speed constraints, memory limitations, and accuracy requirements. MobileBERT delivers the highest performance but is also the largest and slowest. Here is a comparison on common NLP benchmarks:

Model Params Accuracy Latency
MobileBERT 25M 84.3 642 ms
Word2Vec + CNN 1M 79.7 18 ms
FastText + BiLSTM 5M 81.2 51 ms

Latency measured on a Pixel 3 smartphone, accuracy on the IMDb sentiment dataset.

In the rest of this tutorial, we‘ll use MobileBERT since it provides the best balance of performance and efficiency. However, feel free to experiment with different architectures to find the optimal tradeoff for your use case.

Step 3: Training the Model

Now let‘s train our MobileBERT model on the collected dataset. We‘ll use a pretrained checkpoint and fine-tune it for our text classification task.

First, install the required packages:

!pip install tensorflow-text tensorflow-hub

Then load and preprocess the data:

import pandas as pd 
import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text as text

train_df = pd.read_csv(‘train.csv‘)
val_df = pd.read_csv(‘val.csv‘)

label_map = {label: i for i, label in enumerate(train_df[‘label‘].unique())}
train_df[‘label‘] = train_df[‘label‘].map(label_map)
val_df[‘label‘] = val_df[‘label‘].map(label_map)

train_texts = train_df[‘text‘].tolist()  
train_labels = train_df[‘label‘].tolist()
val_texts = val_df[‘text‘].tolist()
val_labels = val_df[‘label‘].tolist()

Load the pre-trained MobileBERT model:

mobilebert_url = "https://tfhub.dev/tensorflow/mobilebert_en_uncased_L-24_H-128_B-512_A-4_F-4_OPT/1"
mobilebert_layer = hub.KerasLayer(mobilebert_url, trainable=True)

Build a Keras model with MobileBERT encoding followed by a dense output layer:

inputs = dict(
    input_word_ids=tf.keras.layers.Input(shape=(None,), dtype=tf.int32),
    input_mask=tf.keras.layers.Input(shape=(None,), dtype=tf.int32),  
    input_type_ids=tf.keras.layers.Input(shape=(None,), dtype=tf.int32))

embedding = mobilebert_layer(inputs)[‘pooled_output‘]
outputs = tf.keras.layers.Dense(len(label_map), activation=‘softmax‘)(embedding)

model = tf.keras.Model(inputs=inputs, outputs=outputs)

Compile the model and start fine-tuning:

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=3e-5),  
              loss=‘sparse_categorical_crossentropy‘,
              metrics=[‘accuracy‘])

model.fit(train_texts, train_labels, 
          validation_data=(val_texts, val_labels),
          epochs=3, batch_size=32)

A few tips for the training process:

  • Use a lower learning rate (e.g. 3e-5) to avoid divergence when fine-tuning pre-trained models
  • Monitor the validation loss and accuracy to check for overfitting
  • Experiment with different hyperparameters like batch size, number of epochs, learning rate schedule
  • Apply techniques like learning rate warmup and discriminative fine-tuning to stabilize training
  • Consider class weights to handle label imbalance

After training, evaluate the final model performance on a held-out test set. TensorBoard is a useful tool for visualizing metrics across runs.

Step 4: Converting to TensorFlow Lite

To deploy our trained MobileBERT model to an Android app, we need to convert it to the TensorFlow Lite format first. This involves serializing the model architecture and weights into a compressed .tflite file.

We use the TFLiteConverter tool:

converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_model = converter.convert()

with open(‘model.tflite‘, ‘wb‘) as f:
    f.write(tflite_model)

Here we apply default optimizations, which enable hybrid quantization (converting weights to 8-bit precision while keeping activations in float32). For even smaller models, you can try full integer quantization, which converts both weights and activations. However, this requires representative data and more careful accuracy tuning.

The model size shrinks from 100 MB to 25 MB after TensorFlow Lite conversion. You can further reduce it through post-training quantization and pruning.

Step 5: Deploying to an Android App

Finally, we integrate the converted model into an Android app for on-device inference. We‘ll build a simple UI where users can input text and view the predicted category.

In Android Studio, create a new project and add the TensorFlow Lite dependencies to your app-level build.gradle file:

dependencies {
    implementation ‘org.tensorflow:tensorflow-lite:2.12.0‘
    implementation ‘org.tensorflow:tensorflow-lite-support:0.5.0‘
    implementation ‘org.tensorflow:tensorflow-lite-metadata:0.5.0‘
}

Copy the model.tflite file to the app‘s assets directory. Then load the model in your main Activity:

import org.tensorflow.lite.support.label.Category;
import org.tensorflow.lite.task.core.BaseOptions;
import org.tensorflow.lite.task.text.nlclassifier.NLClassifier; 

String modelPath = "model.tflite";
BaseOptions options = BaseOptions.builder().build();
NLClassifier classifier = NLClassifier.createFromFileAndOptions(this, modelPath, options);

Set up a TextInputLayout and Button in your activity_main.xml:

<com.google.android.material.textfield.TextInputLayout
    android:id="@+id/textField"
    android:layout_width="match_parent"
    android:layout_height="wrap_content"
    android:hint="Enter text"
    android:layout_margin="16dp">

    <com.google.android.material.textfield.TextInputEditText
        android:layout_width="match_parent"
        android:layout_height="wrap_content" />

</com.google.android.material.textfield.TextInputLayout>

<Button
    android:id="@+id/classifyButton"
    android:layout_width="wrap_content"
    android:layout_height="wrap_content"
    android:layout_margin="16dp"
    android:text="Classify"/>

Then in your MainActivity.java, respond to the button click, run inference on the input text, and display the top predicted category:

import com.google.android.material.textfield.TextInputLayout;

TextInputLayout textField = findViewById(R.id.textField);
Button classifyButton = findViewById(R.id.classifyButton);

classifyButton.setOnClickListener(view -> {
    String inputText = textField.getEditText().getText().toString();
    List<Category> results = classifier.classify(inputText);
    Category topResult = results.get(0);
    String category = topResult.getLabel();
    float confidenceScore = topResult.getScore();

    String message = String.format("Predicted category: %s (%.2f)", category, confidenceScore);
    Toast.makeText(this, message, Toast.LENGTH_LONG).show();
});

Run the app on a device or emulator and test it out with some sample inputs. The model should return sensible predictions with sub-second latency.

Some ways to enhance the user experience:

  • Show a progress bar during inference
  • Color-code results based on the confidence score
  • Add icons or images to represent each category
  • Allow sorting categories by ascending/descending score
  • Support speech input via Android‘s SpeechRecognizer API

Going Beyond

Congratulations on making it this far! Let‘s recap what we‘ve achieved:

  • Collected a labeled text dataset for a custom classification task
  • Trained a state-of-the-art MobileBERT model
  • Converted the model to an optimized TensorFlow Lite format
  • Integrated the model into an interactive Android app

But this is just the beginning – there are many opportunities to extend and improve our system:

  • Model personalization: Adapt the model to each user‘s writing style by tuning on their in-app data
  • Continual learning: Incrementally update the model on new data to prevent distribution shift
  • Multilingual support: Expand to other languages by adding language-specific datasets and models
  • Unsupervised learning: Discover new categories via clustering or topic modeling
  • Active learning: Reduce annotation costs by strategically requesting user labels for uncertain examples
  • Federated learning: Train across decentralized user devices to preserve privacy
  • Model compression: Further shrink the model size through knowledge distillation and network pruning
  • Explainable AI: Provide human-understandable explanations for model predictions

Text classification is a powerful tool with countless use cases, from content moderation and spam filtering to sentiment analysis and intent detection. With the rise of mobile computing, on-device NLP is becoming increasingly important for real-time, privacy-centric and personalized experiences.

We‘ve shown how to combine TensorFlow Lite with state-of-the-art deep learning to bring NLP to the edge. By following the steps in this tutorial, you‘re well equipped to apply these techniques to your own projects. Go forth and build amazing AI-powered apps!

References

Thanks to the TensorFlow, Google and Udacity teams for creating amazing tools and resources to democratize ML!

How useful was this post?

Click on a star to rate it!

Average rating 0 / 5. Vote count: 0

No votes so far! Be the first to rate this post.

You May Like to Read,

Similar Posts