rust-bert

rust-bert 项目介绍

项目概述

rust-bert 是一个使用 Rust 编写的自然语言处理（NLP）库，提供了最先进的自然语言处理模型和处理管道。这个项目是 Hugging Face 的 Transformers 库的移植版，并使用 tch-rs 或 ONNX Runtime 绑定，以及 rust-tokenizers 进行预处理。它支持多线程的词汇标记和 GPU 推理。

主要功能

项目支持的任务包括但不限于：

• 翻译
• 文本摘要
• 多轮对话
• 零样本分类
• 情感分析
• 命名实体识别
• 词性标注
• 问答系统
• 语言生成
• 掩码语言模型
• 句子嵌入
• 关键词提取

通过该库，用户可以较为简单地实现问答、机器翻译、句子生成等复杂的 NLP 任务。

支持的模型

rust-bert 支持广泛的模型架构，如 DistilBERT、BERT、RoBERTa、GPT、T5 等。这些模型各自在支持的任务中表现出色，比如 DistilBERT 和 BERT 在问答系统和命名实体识别中有优异的表现，而 GPT 和 T5 擅长于文本生成和翻译。

部署与安装

rust-bert 依赖于 tch crate，它为 C++ 的 Libtorch API 提供绑定。用户可以手动或自动下载所需的 libtorch 库，并根据系统环境进行配置。

手动安装步骤

1. 从 PyTorch 官方网站 下载 libtorch。
2. 解压下载的库到本地目录。
3. 设置环境变量，以便程序能够正确访问 libtorch。

自动安装

启用 download-libtorch 特性标记即可让脚本自动下载 libtorch 库。默认下载 CPU 版本，但也可以通过设置环境变量 TORCH_CUDA_VERSION 获取 CUDA 版本。

ONNX 支持（可选）

rust-bert 允许通过可选的 onnx 特性启用 ONNX 支持。这需要 ort crate 提供的 onnxruntime C++ 库绑定。

为了使用动态链接，用户需设置 ORT_DYLIB_PATH 来指向下载的 onnxruntime 库的位置。

模型管道

rust-bert 提供了可以直接使用的 NLP 处理管道，其中包括：

• 问答：基于 SQuAD 数据集微调的 DistilBERT 模型。
• 翻译：支持多种源语言和目标语言，利用 Marian 和 M2M100 模型。
• 文本摘要：使用预训练的 BART 模型。
• 对话模型：基于微软 DialoGPT。
• 自然语言生成：可使用 GPT2 和 GPT 进行文本生成。
• 零样本分类：对输入进行零样本分类。
• 情感分析：DistilBERT 模型在 SST-2 数据集上进行微调。
• 命名实体识别：从文本中提取实体。
• 提取关键词：从输入文档中提取关键词和短语。
• 词性标注：对文本进行词法分析。
• 句子嵌入：生成句子向量表示。
• 掩码语言模型：预测句子中的掩码词。

性能与基准

该库与 Python 实现的性能相当，尤其是在计算密集的 NLP 任务中，得益于 Rust 的高效性能管理和并行计算支持。此外，对于文本生成的任务，rust-bert 可能会显著优于 Python，实现更快的处理速度。

预训练和自定义模型

rust-bert 也允许加载预训练模型，并对用户自定义的 Transformer 模型进行支持。用户可以使用本地模型文件或远程资源进行模型加载。

致谢

感谢 Hugging Face 提供与此 Rust 库兼容的模型权重，使用户可以快速使用预训练模型解决现实中的 NLP 问题。

rust-bert 是一个功能强大的 Rust 库，提供了直观的接口和强大的 NLP 能力，无论是开发者还是 NLP 研究人员都可以从中获益。