Text Classifier

MediaPipe Text Classifier task lets you classify text into a set of defined categories, such as positive or negative sentiment. The categories are defined during the training of the model. This task operates on text data with a machine learning (ML) model as static data and outputs a list of categories and their likelihood scores.

Interactive demo Text Classifier Classify sentiment and tone of natural-language text in sub-millisecond latency. Launch interactive demo Open in Web DemosEdit in StackBlitz On-device benchmark · Pixel 6 Average Word Embedding CPU0.14ms

Open in Web DemosEdit in StackBlitz

Quickstart

Install the SDK for your platform, initialize the task with a model, and run inference.

import { FilesetResolver, TextClassifier } from "@mediapipe/tasks-text";

const MODEL_URL = "https://storage.googleapis.com/mediapipe-models/text_classifier/bert_classifier/float32/latest/bert_classifier.tflite";

const textFiles = await FilesetResolver.forTextTasks(
  "https://cdn.jsdelivr.net/npm/@mediapipe/tasks-text@latest/wasm"
);
const textClassifier = await TextClassifier.createFromOptions(textFiles, {
  baseOptions: { modelAssetPath: MODEL_URL },
  maxResults: 5,
});

const inputText = "An unforgettable and genuinely moving film!";
const result = textClassifier.classify(inputText);

import com.google.mediapipe.tasks.core.BaseOptions
import com.google.mediapipe.tasks.text.textclassifier.TextClassifier
import com.google.mediapipe.tasks.text.textclassifier.TextClassifier.TextClassifierOptions

val baseOptions = BaseOptions.builder()
    .setModelAssetPath("bert_classifier.tflite")
    .build()
val options = TextClassifierOptions.builder()
    .setBaseOptions(baseOptions)
    .build()

val textClassifier = TextClassifier.createFromOptions(context, options)

val inputText = "An unforgettable and genuinely moving film!"
val result = textClassifier.classify(inputText)

import com.google.mediapipe.tasks.core.BaseOptions;
import com.google.mediapipe.tasks.text.textclassifier.TextClassifier;
import com.google.mediapipe.tasks.text.textclassifier.TextClassifier.TextClassifierOptions;
import com.google.mediapipe.tasks.text.textclassifier.TextClassifierResult;

BaseOptions baseOptions = BaseOptions.builder()
    .setModelAssetPath("bert_classifier.tflite")
    .build();
TextClassifierOptions options = TextClassifierOptions.builder()
    .setBaseOptions(baseOptions)
    .build();

TextClassifier textClassifier = TextClassifier.createFromOptions(context, options);

String inputText = "An unforgettable and genuinely moving film!";
TextClassifierResult result = textClassifier.classify(inputText);

import Foundation
import MediaPipeTasksText

guard let modelPath = Bundle.main.path(forResource: "bert_classifier", ofType: "tflite") else { return }

let options = TextClassifierOptions()
options.baseOptions.modelAssetPath = modelPath
options.scoreThreshold = 0.6

let textClassifier = try TextClassifier(options: options)

let inputText = "An unforgettable and genuinely moving film!"
let result = try textClassifier.classify(text: inputText)

import mediapipe as mp

BaseOptions = mp.tasks.BaseOptions
TextClassifier = mp.tasks.text.TextClassifier
TextClassifierOptions = mp.tasks.text.TextClassifierOptions

model_path = "bert_classifier.tflite"
options = TextClassifierOptions(base_options=BaseOptions(model_asset_path=model_path))

input_text = "An unforgettable and genuinely moving film!"
with TextClassifier.create_from_options(options) as classifier:
    classification_result = classifier.classify(input_text)

Task details

This section describes the capabilities, inputs, outputs, and configuration options of this task.

Features

• Input text processing - Support out-of-graph tokenization for models without in-graph tokenization
• Multiple classification heads - Each head can use its own category set
• Label map locale - Set the language used for display names
• Score threshold - Filter results based on prediction scores
• Top-k classification results - Limit the number of classification results
• Label allowlist and denylist - Specify the categories detected

Configuration options

This task has the following configuration options:

Task output

Each TextClassifierResult contains one Classifications entry per classification head:

Example TextClassifierResult payload

The following example shows the sentiment classification output for the input text "an imperfect but overall entertaining mystery":

{
  classifications: [
    {
      headIndex: 0,
      categories: [
        { index: 0, score: 0.8904, categoryName: "positive" },
        { index: 1, score: 0.1096, categoryName: "negative" }
      ]
    }
  ]
}

Models

Text Classifier can be used with more than one ML model. Start with the default, recommended model for your target platform when you start developing with this task. The other available models typically make trade-offs between performance, accuracy, resolution, and resource requirements, and in some cases, include additional features.

The pretrained models are trained for sentiment analysis, and predict whether the input text's sentiment is positive or negative. The models were trained on the SST-2 (Stanford Sentiment Treebank) dataset, which consists of movie reviews labeled as either positive or negative. Note that the models only support English. Since they were trained on a dataset of movie reviews, you may see reduced quality for text covering other topic areas.

BERT-classifier model (recommended)

This model uses a BERT-based architecture (specifically, the MobileBERT model) and is recommended because of its high accuracy. It contains metadata that allows the task to perform out-of-graph BERT tokenization.

Average word embedding model

This model uses an average word-embedding architecture. This model offers a smaller model size and lower latency at the cost of a lower prediction accuracy compared to the BERT-classifier. Customizing this model through additional training is also faster than doing training of the BERT-based classifier. This model contains metadata that allows the task to perform out-of-graph regex tokenization.

Task benchmarks

Here are the task benchmarks for the whole pipeline based on the preceding pre-trained models. The latency result is the average latency on Pixel 6 using CPU / GPU.