集成 BERT 自然语言分类器

任务库 BertNLClassifier API 与 NLClassifier 非常相似,后者将输入文本分类到不同的类别中,但此 API 专为与 Bert 相关的模型量身定制,这些模型需要在 TFLite 模型之外进行 Wordpiece 和 Sentencepiece 分词。

BertNLClassifier API 的主要功能

  • 将单个字符串作为输入,使用该字符串进行分类,并输出对作为分类结果的配对。

  • 在输入文本上执行图外 WordpieceSentencepiece 分词。

支持的 BertNLClassifier 模型

以下模型与 BertNLClassifier API 兼容。

在 Java 中运行推理

步骤 1:导入 Gradle 依赖项和其他设置

.tflite 模型文件复制到将运行模型的 Android 模块的 assets 目录。指定该文件不应压缩,并将 TensorFlow Lite 库添加到模块的 build.gradle 文件中

android {
    // Other settings

    // Specify tflite file should not be compressed for the app apk
    aaptOptions {
        noCompress "tflite"
    }

}

dependencies {
    // Other dependencies

    // Import the Task Text Library dependency (NNAPI is included)
    implementation 'org.tensorflow:tensorflow-lite-task-text:0.4.4'
}

步骤 2:使用 API 运行推理

// Initialization
BertNLClassifierOptions options =
    BertNLClassifierOptions.builder()
        .setBaseOptions(BaseOptions.builder().setNumThreads(4).build())
        .build();
BertNLClassifier classifier =
    BertNLClassifier.createFromFileAndOptions(context, modelFile, options);

// Run inference
List<Category> results = classifier.classify(input);

有关更多详细信息,请参阅 源代码

在 Swift 中运行推理

步骤 1:导入 CocoaPods

在 Podfile 中添加 TensorFlowLiteTaskText pod

target 'MySwiftAppWithTaskAPI' do
  use_frameworks!
  pod 'TensorFlowLiteTaskText', '~> 0.4.4'
end

步骤 2:使用 API 运行推理

// Initialization
let bertNLClassifier = TFLBertNLClassifier.bertNLClassifier(
      modelPath: bertModelPath)

// Run inference
let categories = bertNLClassifier.classify(text: input)

有关更多详细信息,请参阅 源代码

在 C++ 中运行推理

// Initialization
BertNLClassifierOptions options;
options.mutable_base_options()->mutable_model_file()->set_file_name(model_path);
std::unique_ptr<BertNLClassifier> classifier = BertNLClassifier::CreateFromOptions(options).value();

// Run inference with your input, `input_text`.
std::vector<core::Category> categories = classifier->Classify(input_text);

有关更多详细信息,请参阅 源代码

在 Python 中运行推理

步骤 1:安装 pip 包

pip install tflite-support

步骤 2:使用模型

# Imports
from tflite_support.task import text

# Initialization
classifier = text.BertNLClassifier.create_from_file(model_path)

# Run inference
text_classification_result = classifier.classify(text)

有关配置 BertNLClassifier 的更多选项,请参阅 源代码

示例结果

以下是用 Model Maker 中的 MobileBert 模型对电影评论进行分类的结果示例。

输入:"it's a charming and often affecting journey"

输出

category[0]: 'negative' : '0.00006'
category[1]: 'positive' : '0.99994'

使用您自己的模型和测试数据试用简单的 BertNLClassifier 的 CLI 演示工具

模型兼容性要求

BetNLClassifier API 预计 TFLite 模型具有强制性的 TFLite 模型元数据

元数据应满足以下要求

  • 用于 Wordpiece/Sentencepiece 分词器的 input_process_units

  • 3 个输入张量,名称分别为 "ids"、"mask" 和 "segment_ids",用于分词器的输出

  • 一个类型为 float32 的输出张量,可以选择附加标签文件。如果附加了标签文件,该文件应为纯文本文件,每行一个标签,标签数量应与模型输出的类别数量匹配。