“
「 Rust 与 LLM」 是本合集的主题系列之一,本文为正文第一篇。阅读本系列文章不需要有数学基础。
本合集将优先完成该主题系列文章,所以其他主题的文章优先级将降低。
「 Rust 与 LLM」主题系列将专注于自然语言处理、 Transfomer 架构和大模型相关内容,依托 Rust 开源生态和 HuggingFace 的相关 Rust 库,探秘从模型训练到模型部署、模型量化与 WebAssembly 轻量化部署的技术原理。
正如本主题系列第一篇文章所说,自从 BERT 和 GPT 模型取得重大成功之后, Transformer 结构已经替代了循环神经网络 (RNN) 和卷积神经网络 (CNN),成为了当前大语言模型的主流。
所以我们的系列主题就从自然语言处理开始,逐步入门 Transfomer 架构大语言模型。
自然语言处理(Natural Language Processing,NLP)是人工智能(AI)和计算语言学领域的一个分支,它旨在使计算机能够理解、解释和生成人类语言的数据。NLP涵盖了从基础的文本和语音处理到复杂的对话系统、情感分析、语言翻译、以及语义理解等广泛的应用。
自然语言 Transfomer 架构也是一种神经网络。由于神经网络模型不能直接处理文本,因此需要先将文本转换为特殊的数字,这个过程被称为**编码 (Encoding)**,其包含两个步骤:
使用分词器 (tokenizer) 将文本按词、子词、字符切分为词元( tokens);
将所有的词元映射到对应的 token ID。
本文的主角 tokenizers 就是用来处理这个文本编码过程。
分词器(Tokenizer)是 NLP 中的一个基础组件,它的任务是将文本分解成更小的单元,通常是单词、短语或者句子。分词(亦称标记化)是 NLP 预处理步骤中的关键环节,因为大多数语言模型和算法都是在单词或短语级别上操作的。准确的分词直接影响到后续任务的效果,如情感分析、实体识别、语法分析等。
分词器的主要类型:
基于规则的分词器:这种分词器依据预定义的规则(如空格、标点符号等)来切分文本。它简单直观,但可能无法很好地处理复合词、缩写词等特殊情况。
基于统计的分词器:利用机器学习算法,根据大量文本数据学习单词的边界。它能更灵活地处理语言的多样性,但需要大量的训练数据。
基于深度学习的分词器:使用深度神经网络,如循环神经网络(RNN)或Transformer,来学习文本中的单词边界。这种方法可以更准确地识别复杂的语言结构,但计算成本较高。
我们将要学习的 HuggingFace 出品的分词器 tokeninzers 则属于基于统计的分词器类别。
一个基本的基于统计的分词器框架,包含以下主要部分:
初始化:定义了词汇表、合并规则、特殊标记等基本属性。
训练:应当根据输入文本和指定的词汇表大小来训练 BPE 分词器。
编码与解码:分别用于将文本编码为整数序列和将整数序列解码回文本。
构建词汇表:根据合并规则动态构建词汇表,其中包括基于字节的原始词汇表和通过合并得到的标记。
保存与加载:允许将训练好的分词器规则保存到文件中,并从文件中加载规则,方便分词器的复用和分发。
主要是因为不同的语言、领域或应用场景下,文本的结构和用词习惯可能有很大差异。训练分词器可以帮助模型更好地理解和处理特定于任务或数据集的语言特性,从而提高自然语言处理(NLP)任务的性能。
以下是几个具体原因:
不同语言的结构差异:例如,中文是一种以字为基本单位的语言,没有明显的单词分隔符,而英文单词之间通常由空格分隔。训练分词器可以帮助模型学习到针对特定语言的最佳分词策略。
处理复合词和派生词:在许多语言中,新词、复合词和派生词的形成非常活跃。通过训练,分词器可以更好地识别和处理这些词汇,避免将它们错误地分割成未知或不相关的部分。
构建高效词汇表:通过训练,可以构建一个既包含高频词汇又能有效覆盖训练数据的词汇表,从而在保持模型性能的同时减少其大小。
处理未见词(OOV)问题:训练分词器可以通过学习子词单元(如BPE算法)来更好地处理训练数据中未出现的词汇,从而提高模型对新词的适应能力。
领域专有名词和术语:特定领域(如医学、法律或技术)常常有自己独特的术语和表达方式。训练分词器可以帮助模型学习这些领域特定的词汇和用法,提高在特定领域内的准确性和效率。
适应语言变化:语言是不断发展变化的,尤其是在互联网和社交媒体的文本中,新词汇、网络用语和表情符号等频繁出现。通过针对这些数据训练分词器,可以使模型更好地理解和处理这类文本。
减少错误率:通过训练,分词器可以更准确地识别和分割词汇,减少分词错误,从而为下游任务提供更准确的输入。
提升下游任务效果:无论是文本分类、情感分析、机器翻译还是问答系统,准确的分词都是提高模型性能的关键。训练分词器可以优化模型在特定任务上的表现。
tokenizers[1] 是用 Rust 语言实现的高性能分词器,主要被用在 HuggingFace 的 transfomer pipeline中。本文将拆解该分词器架构和主要逻辑。
tokenizers 的基本逻辑结构如下图。
该逻辑结构代表了分词器的工作流:文本规范化(Normalization)、预分词(Pre-tokenization)、分词(Tokenization)、以及后处理(Post-processing)。
tokenizers 的代码组织结构和架构也是遵循该逻辑而设计的。
在 tokenizers/
源代码进一步细分为多个模块,包括:
src/normalizers/: 包含文本规范化逻辑,例如去除空格、小写化等。
src/pre_tokenizers/: 包含预分词逻辑,例如按空格、标点符号等分割。
src/models/: 包含不同分词模型的实现,例如 BPE(Byte Pair Encoding)、Unigram、WordLevel、WordPiece 等。
src/processors/: 包含后处理逻辑,例如添加特殊标记(如 [CLS]、[SEP])等。
src/tokenizer/: 包含将上述组件综合起来的分词器逻辑。
src/utils/: 包含各种实用工具和功能,例如缓存、并行处理等。
binding/
目录下则提供了 Python 语言(基于 pyo3
)和 NodeJS 绑定(基于 npi
),据说 Ruby 语言绑定正在进行中。
其中 tokenizer 为上层模块,它提供了该库的公共接口和类型等。在下面的组件模块中会用到上层模块定义的一些 trait 和类型。
normalizers 模块提供的 Normalizer
负责对输入字符串进行预处理,以便根据给定的用例对其进行规范化。一些常见的规范化示例包括 Unicode 规范化算法(NFD,NFKD,NFC 和 NFKC),小写等等... tokenizers
的特殊之处在于会在规范化过程中保持对齐的跟踪。这对于将生成的标记映射回输入文本是至关重要的。
Unicode 规范化标准[2] 中定义了 NFD、NFC、NFKD 和 NFKC 四种规范化形式。
NFD(Normalization Form Decomposed),将字符分解为基础字符和修饰符号。例如,将带重音的字符(如 "é")分解为基础字符 "e" 和重音符号。
NFC(Normalization Form Composed),将分解的字符重新组合为单一字符(如果可能)。例如,将基础字符 "e" 和重音符号组合成 "é"。
NFKD(Normalization Form Compatibility Decomposed),进行兼容性分解,将字符分解为与之兼容的基础字符,可能会改变字符的视觉表现。例如,将 "fi"(连字 fi)分解为两个单独的字符 "f" 和 "i"。
NFKC(Normalization Form Compatibility Composed),进行兼容性分解后,尽可能地将分解的字符重新组合。
总的来说,Unicode 标准定义了两种字符之间的正式等价关系:规范等价(NFD/NFC) 和 **兼容等价(NFKD/NFKC)**。规范等价是指字符或字符序列之间的基本等价关系,它们表示相同的抽象字符,并且在正确显示时应该始终具有相同的视觉外观和行为。兼容等价性是字符或字符序列之间的一种较弱的等价关系,它们表示相同的抽象字符(或抽象字符序列),但可能具有不同的视觉外观或行为。兼容等价形式的视觉外观通常构成了字符(或字符序列)的预期视觉外观范围的子集。然而,这些变体形式可能在某些文本环境中表示一个重要的视觉区别,但在其他环境中则不重要。
除了 Unicode 规范化之外,还包括下面规范化形式:
小写字母(Lowercase ),将所有大写字母替换为小写字母。
剥离(Strip),删除输入的指定边(左边、右边或两边)的所有空白字符。
去除重音(StripAccents),删除 Unicode 中的所有重音符号(与NFD一起使用以保持一致性)。
替换(Replace),替换自定义的字符串或正则表达式,并将其更改为给定的内容。
Bert 规范化器(BertNormalizer),提供了原始BERT中使用的规范化器的实现。
序列(Sequence),用于组合多个规范化器,按照提供的顺序运行。
关键架构代码
在 公共接口模块 src/tokenizer
中定义了规范化 trait 和 NormalizedString 结构体:
`// src/tokenizer/mod.rs /// Takes care of pre-processing strings. pub trait Normalizer { fn normalize(&self, normalized: &mut NormalizedString) -> Result<()>; } // src/tokenizer/normalizer.rs #[derive(Default, Debug, Clone, PartialEq, Eq)] pub struct NormalizedString { /// The original version of the string, before any modification original: String, /// The normalized version of the string, after all modifications normalized: String, /// Mapping from normalized string to original one: (start, end) for each /// byte of the normalized string // 包含了对齐信息,这是后期矩阵运算要求 alignments: Vec<(usize, usize)>, /// If this NormalizedString is a slice of a bigger one, we keep the track /// of the missing part, so that we can still give offsets from this original /// string. original_shift: usize, } `
这是对规范化文化的统一抽象。每个字段的意义注释已经很清楚了。NormalizedString 提供了一些基本操作,其中重要的操作是 transform_range
方法,用于处理对齐,为后期矩阵运算提供支持。
然后为 NormalizerWrapper 枚举实现了 Normalizer trait,这算是一种静态分发,相比于使用 trait 对象动态分发,性能更好,接近零成本抽象。
`// src/normalizers/mod.rs pub enum NormalizerWrapper { BertNormalizer(BertNormalizer), StripNormalizer(Strip), StripAccents(StripAccents), NFC(NFC), NFD(NFD), NFKC(NFKC), NFKD(NFKD), Sequence(Sequence), Lowercase(Lowercase), Nmt(Nmt), Precompiled(Precompiled), Replace(Replace), Prepend(Prepend), } // src/normalizers/mod.rs impl Normalizer for NormalizerWrapper { fn normalize(&self, normalized: &mut NormalizedString) -> crate::Result<()> { match self { Self::BertNormalizer(bn) => bn.normalize(normalized), Self::StripNormalizer(sn) => sn.normalize(normalized), Self::StripAccents(sn) => sn.normalize(normalized), Self::NFC(nfc) => nfc.normalize(normalized), Self::NFD(nfd) => nfd.normalize(normalized), Self::NFKC(nfkc) => nfkc.normalize(normalized), Self::NFKD(nfkd) => nfkd.normalize(normalized), Self::Sequence(sequence) => sequence.normalize(normalized), Self::Lowercase(lc) => lc.normalize(normalized), Self::Nmt(lc) => lc.normalize(normalized), Self::Precompiled(lc) => lc.normalize(normalized), Self::Replace(lc) => lc.normalize(normalized), Self::Prepend(lc) => lc.normalize(normalized), } } } `
这里手工枚举实现静态分发的前提是,规范化形式的种类是已知和有限的,并且都是库内部提供的。
其内部使用的 Unicode 规范化方法是由第三方库 `unicode_normalization_alignments`[3] 提供的。
另外值得看的是 BertNormalizer :
``#[derive(Copy, Clone, Debug, Deserialize, Serialize)] #[serde(tag = "type")] #[non_exhaustive] pub struct BertNormalizer { /// Whether to do the bert basic cleaning: /// 1. Remove any control characters /// 2. Replace all sorts of whitespace by the classic one ` ` pub clean_text: bool, /// Whether to put spaces around chinese characters so they get split pub handle_chinese_chars: bool, /// Whether to strip accents pub strip_accents: Option<bool>, /// Whether to lowercase the input pub lowercase: bool, } ``
它对应的是 Bert 模型(Google 出的 transformer 架构模型),它有专门的规范化要求:
首先要清理各种特殊字符,比如移除控制符,将所有空白字符转为空格,后面预分词阶段会按空格进行预分词。
判断是否为中文,是的话在前后添加空格,否则原样输出。
如果需要,则去除重音。
如果需要,将所有字符都转为小写。
在 util.rs
中提供了 Sequence
:
`pub struct Sequence { normalizers: Vec<NormalizerWrapper>, } impl Normalizer for Sequence { fn normalize(&self, normalized: &mut NormalizedString) -> Result<()> { for normalizer in &self.normalizers { normalizer.normalize(normalized)?; } Ok(()) } } `
可以迭代处理规范化器。
其他规范化模块都是比较常规的操作了,这里不做多讲。
pre_tokenizers 模块负责根据一组规则对输入进行预分词。例如,如果你不希望在标记内部有空格,则可以使用 PreTokenizer
在这些空格上进行拆分。
预处理操作具体包括:
ByteLevel,在将所有字节重新映射为一组可见字符的同时,根据空格进行分割。这种技术由 OpenAI 与 GPT-2 引入,并具有一些或多或少的良好特性。由于它映射到字节,使用这种方式的分词器只需要256个字符作为初始字母表(一个字节可以有的值的数量),而不是130,000多个Unicode字符。这样就做到可以用256个 token 来表示任何东西。
Whitespace,按空格分割单词。
WhitespaceSplit,按任何空格进行分割,包括单词内部。
Punctuation,将隔离所有标点符号字符。
Metaspace,在空格处分割并用特殊字符“▁”(U+2581)替换它们。
CharDelimiterSplit,根据给定的字符进行分割。
Digits,按数字进行分割。
Split,多功能的预分词器,根据提供的模式和行为进行分割。如果需要,可以反转模式。模式应该是自定义字符串或正则表达式。
Sequence,允许组合多个 PreTokenizer
,按照给定的顺序运行。
关键代码架构:
在公共模块 src/tokenizer/mod.rs
中提供了 PreTokenizer
trait,它类似于 Normalizer
trait 的抽象。
PreTokenizer
负责执行预分割步骤。它将给定的字符串分割成多个子字符串,并跟踪这些子字符串与 NormalizedString
之间的偏移量。在某些情况下, PreTokenizer
可能需要修改给定的 NormalizedString
,以确保我们可以完全跟踪偏移量和与原始字符串的映射关系。
``// src/tokenizer/mod.rs pub trait PreTokenizer { fn pre_tokenize(&self, pretokenized: &mut PreTokenizedString) -> Result<()>; } // src/tokenizer/pre_tokenizer.rs pub struct PreTokenizedString { original: String, splits: Vec<Split>, } pub struct Split { /// The underlying `NormalizedString`. Each SubString is represented by a `NormalizedString` /// and in the end we might be carrying a lot of SubString representing various parts of the /// original input string. normalized: NormalizedString, /// Optional Tokens associated to this Split tokens: Option<Vec<Token>>, } // src/tokenizer/pattern.rs /// Pattern used to split a NormalizedString pub trait Pattern { /// Slice the given string in a list of pattern match positions, with /// a boolean indicating whether this is a match or not. /// /// This method *must* cover the whole string in its outputs, with /// contiguous ordered slices. fn find_matches(&self, inside: &str) -> Result<Vec<(Offsets, bool)>>; } ``
PreTokenizedString
结构体也实现了很多操作,这里不一一罗列。
这里面值得看的模块是 byte_level.rs
。其核心功能:
字节到字符的转换:bytes_char
函数创建了一个从字节(u8)到 Unicode 字符(char)的映射。这个映射包括了 ASCII 可打印字符范围内的所有字符,以及扩展的字节范围。这对于处理基于字节级别的文本编码特别有用。
正则表达式匹配:RE
静态变量定义了一个正则表达式,用于匹配一个 token。这个正则表达式涵盖了英语的所有权和缩略形式,以及字母、数字和非空格非字母数字字符。这有助于在分词过程中识别和分割单词和标点符号。
字节级分词器:ByteLevel
结构体实现了 PreTokenizer
、Decoder
和 PostProcessor
trait,提供了在字节级别处理 BPE 分词(下一篇细说)所需的所有步骤。这包括将 Unicode 字符转换为它们的字节级表示,使用正则表达式进行分割,以及在分词后将字节转换回字符。
代码细节就不在这里展示了。整体而言,这是 Rust 实现的一个字节级的预分词器,该预分词器适用于需要在最细粒度上处理文本的 NLP 任务,例如使用 BPE 算法的语言模型训练。通过对 Unicode 字符进行字节级处理,该实现能够有效支持多语言文本并提高模型对未见词的泛化能力。
这种方法之所以能够使用仅 256 个 token 表示任何内容,核心原因在于它直接基于字节(byte)级别的处理,而不是更高层次的 Unicode 字符。在这种情况下,分词器的初始字母表实际上是所有可能的字节值(0-255),这恰好对应于一个字节能表示的所有可能值。通过这种方式,分词器能够覆盖所有可能的输入文本,因为任何文本文件最终都是以字节序列的形式存储的。这里的关键在于两个转换过程:
字节到字符的映射。通过将每个字节映射到一个特定的 Unicode 字符,我们实际上创建了一种能够通过 256 个基本单元(即字节值)来表示任何文本的系统。这种映射使得原始的字节序列可以被视为一串 Unicode 字符,进而可以用于后续的文本处理和分词。
分词器处理字节序列。当处理文本时,分词器不是直接处理高层次的 Unicode 字符,而是将文本转换成字节序列,并在这个序列上进行操作。这意味着,无论原始文本使用什么语言或符号,它都被归一化为一个统一的字节序列表。因此,使用 256 个基本 token(对应于256个可能的字节值)就足以表示任何文本内容。
分词模型是实际进行分词的核心算法,因此,它是分词器的唯一必需组件,其他规范化、预分词其实都是可选的。
在 tokenizers 库的 model 模块则包含了不同的分词模型,它们对应了不同的分词方式。
BPE(Byte Pair Encoding)
原理:BPE 通过统计字符对(byte pairs)的频率并迭代地将最常见的字符对合并成新的符号,从而实现文本的分词。这个过程重复进行,直到达到预定的词汇表大小或没有更多的合并可以进行。
优点:能够有效处理未知词或罕见词,因为它可以将这些词分解成较小的、已知的子词单元。
应用:在一些预训练语言模型中广泛使用,如GPT系列。
Unigram
原理:Unigram 模型基于单词出现的概率来进行分词。它首先假设文本中的每个片段(可以是字符、子词等)都是独立选择的,然后通过迭代优化来确定最终的词汇表,每次迭代都会尝试删除概率最低的词汇,直到达到设定的词汇表大小。
优点:自动学习词汇表,能够平衡分词的粒度,同时处理未见过的词。
应用:在一些最新的语言模型和机器翻译系统中被采用。
WordLevel
原理:WordLevel 分词模型基于空格和标点符号直接将文本分割成词汇。它通常使用一个固定的词汇表,只保留训练数据中出现频率较高的单词。
优点:简单直观,易于实现和理解。
局限:无法有效处理未知词,对于拼写错误或变体形式也较为敏感。
WordPiece
原理:WordPiece 类似于 BPE,但在合并过程中引入了一个选择最优分割的方法。它以贪心算法逐步构建词汇表,每次迭代选择能最大化语料库似然的合并。
优点:与 BPE 类似,能够有效处理未知词,同时优化分词的选择过程。
应用:被广泛用于 Google 的 BERT 以及其他预训练语言模型中。
关于分词模型的原理和实现,我们放到下一篇内容细讲。
整个流程完成后,有时会希望在将分词字符串输入模型之前插入一些特殊标记,例如,在 Bert 模型中,“[CLS] Rust 语言很棒 [SEP]”。
processors 模块就是用来处理这样的需求。
关键代码架构:
分词最终输出的是 Encoding :
``/// Represents the output of a `Tokenizer`. #[derive(Default, PartialEq, Debug, Clone, Serialize, Deserialize)] pub struct Encoding { /// IDs produced by the `Tokenizer` ids: Vec<u32>, /// Type of the IDs type_ids: Vec<u32>, /// Tokens associated to each ID tokens: Vec<String>, /// Indice of the word associated to each token/ID words: Vec<Option<u32>>, /// Offsets of the token/ID from the NormalizedString offsets: Vec<Offsets>, /// Mask identifying special tokens special_tokens_mask: Vec<u32>, /// Mask identifying padding tokens for the attention mechanism attention_mask: Vec<u32>, /// A list of overflowing Encoding generated when we got truncated overflowing: Vec<Encoding>, /// Ranges of tokens covered by each sequence. If this is empty we consider /// there is only one sequence in this Encoding, and that it covers the entire range. sequence_ranges: HashMap<usize, Range<usize>>, } ``
字段解释:
ids
:分词器生成的每个 token 的 ID。
type_ids
:每个 ID 的类型,用于区分不同类型的 token,如在特定任务中区分句子A和句子B的 token。
tokens
:与每个 ID 关联的 token 字符串。
words
:与每个 token/ID 关联的单词索引,为 Option<u32>
类型,允许某些 token 没有直接关联的单词。
offsets
:每个 token/ID 在原始文本中的偏移量(起始位置和结束位置)。
special_tokens_mask
:标记特殊 token 的掩码,用于识别如 [CLS]
、[SEP]
等特殊 token。
attention_mask
:用于注意力机制的掩码,标记哪些 token 应该被模型注意。
overflowing
:当文本被截断时,生成的溢出 Encoding
的列表。
sequence_ranges
:由序列索引到其覆盖的 token 范围的映射,用于处理包含多个序列的情况。
在 src/tokenizer/mod.rs
中也定义了 PostProcessor
trait 来抽象后处理操作,因为不同的模型后处理行为不同,这里需要通过 trait 来统一接口。
``/// A `PostProcessor` has the responsibility to post process an encoded output of the `Tokenizer`. /// It adds any special tokens that a language model would require. pub trait PostProcessor { /// Returns the number of tokens that will be added during the processing step fn added_tokens(&self, is_pair: bool) -> usize; /// Process both encodings and returns a new merged one fn process( &self, encoding: Encoding, pair_encoding: Option<Encoding>, add_special_tokens: bool, ) -> Result<Encoding> { let mut encodings = if let Some(pair_encoding) = pair_encoding { vec![encoding, pair_encoding] } else { vec![encoding] }; encodings.iter_mut().enumerate().for_each(|(i, encoding)| { encoding.set_sequence_id(i); encoding .get_overflowing_mut() .iter_mut() .for_each(|encoding| encoding.set_sequence_id(i)); encoding.set_type_ids(vec![i as u32; encoding.len()]); }); let encodings = self.process_encodings(encodings, add_special_tokens)?; Ok(Encoding::merge(encodings, false)) } /// Process any amount of encodings and returns a series of encoding (might merge them) fn process_encodings( &self, encodings: Vec<Encoding>, add_special_tokens: bool, ) -> Result<Vec<Encoding>>; } impl dyn PostProcessor { pub fn default_process( encodings: Vec<Encoding>, _add_special_tokens: bool, ) -> Result<Vec<Encoding>> { match encodings.len() { 1 => Ok(encodings), _ => { let mut final_encoding = Encoding::default(); for (i, mut encoding) in encodings.into_iter().enumerate() { encoding.set_sequence_id(i); final_encoding.merge_with(encoding, false); } Ok(vec![final_encoding]) } } } } ``
并且这里使用了 trait 对象,因为后处理操作不是有限的,所以无法静态分发了,但是为 trait 对象实现一个默认的后处理操作是可能的。
拿 Bert 模型对应的后处理来说,基本代码结构如下:
`#[derive(Serialize, Deserialize, Clone, Debug, PartialEq, Eq)] #[serde(tag = "type")] pub struct BertProcessing { sep: (String, u32), cls: (String, u32), } impl PostProcessor for BertProcessing { fn added_tokens(&self, is_pair: bool) -> usize { if is_pair { 3 } else { 2 } } fn process_encodings( &self, mut encodings: Vec<Encoding>, add_special_tokens: bool, ) -> Result<Vec<Encoding>> { // ...... 此处省略 } } `
解码器知道如何从分词器使用的 ID 返回可读的文本。例如,一些 Normalizer
和 PreTokenizer
使用特殊字符或标识符,需要进行还原。
同样,这块内容放到下一篇结合 BPE 算法来讲。
在 src/utils
模块中定义了分词器相关的优化策略。包括:
cache.rs
,用于加速 BPE 分词算法的多线程缓存。
iter.rs
,用于高效处理迭代器中的 Result<T, E>
项,在迭代时不需要多余的内存拷贝,而主要涉及到迭代器状态的检查和更新。以及可以读取具有行尾的文本行的类似于std::io::BufRead
的 lines
方法。
parallelism.rs
,使用了 rayon
库来实现并行迭代器,并引入了 rayon_cond
的条件迭代器(CondIterator
)。它允许根据环境变量 TOKENIZERS_PARALLELISM
或其他条件来决定迭代器是以并行模式还是串行模式执行。在需要根据运行时条件动态选择最优执行策略的场景中非常有用。这套封装允许开发者编写灵活的并行/串行代码,根据运行时条件动态调整执行策略,从而在保持代码简洁的同时,最大化性能表现。特别是在需要处理大量数据且执行时间敏感的场景下,如 NLP 任务处理、数据分析等,根据可用资源智能选择并行或串行处理可以显著提高程序的效率和响应速度。
在 paralleism.rs
中包含了三个 trait:
MaybeParallelIterator
trait,该 trait 定义了将任何迭代器转换为可以根据条件并行或串行执行的迭代器的能力。它定义了两个方法:
into_maybe_par_iter
根据环境变量自动选择并行或串行执行。
into_maybe_par_iter_cond
根据环境变量和额外的布尔条件选择执行方式。
MaybeParallelBridge
trait,该 trait 允许将序列迭代器转换为 CondIterator
,这种迭代器可以根据环境变量和条件以并行或串行方式执行。
MaybeParallelSlice
trait,该 trait 为切片(slice)提供了类似的功能,允许根据条件将切片的 chunks
操作转换为可以并行或串行执行的迭代器。
tokenizers 分词库在使用批量 encoding(encode_batch
) 时采用了并行编码策略。
案例来自 HuggingFace 官方有训练分词器的指南[4]。结合前面讲的流程和代码架构理解一下。
``use tokenizers::models::bpe::BPE; use tokenizers::models::bpe::BpeTrainer; use tokenizers::pre_tokenizers::whitespace::Whitespace; use tokenizers::processors::template::TemplateProcessing; use tokenizers::PaddingParams; fn main() -> Result<(),Error> { // 使用 BPE 分词模型实例化 let mut tokenizer: TokenizerImpl< BPE, NormalizerWrapper, PreTokenizerWrapper, PostProcessorWrapper, DecoderWrapper, > = TokenizerImpl::new( BPE::builder() .unk_token("[UNK]".to_string()) .build() .unwrap(), ); // 通过 BpeTrainer 实例化一个训练器 // 特殊标记列表的顺序很重要: // 这里 `"[UNK]"` 将获得ID 0, `"[CLS]"` 将获得ID 1,依此类推。 let mut trainer = BpeTrainer::builder() // 在后面插入词汇表的特殊 token .special_tokens(vec![ AddedToken::from("[UNK]", true), AddedToken::from("[CLS]", true), AddedToken::from("[SEP]", true), AddedToken::from("[PAD]", true), AddedToken::from("[MASK]", true), ]) .build(); // 到此为止,可以训练分词器了,但这并不是最佳选择 // 如果没有一个预分词器将输入分割成单词,可能会得到一些重叠多个单词的 token // 这里想训练一个子词BPE分词器,并且使用最简单的预分词器:按空格分割。 tokenizer.with_pre_tokenizer(Whitespace {}); // 只需使用任何想要使用的文件列表来调用 `Tokenizer.train` 方法 let files = vec![ "data/wikitext-103-raw/wiki.train.raw".into(), "data/wikitext-103-raw/wiki.test.raw".into(), "data/wikitext-103-raw/wiki.valid.raw".into(), ]; tokenizer.train_from_files(&mut trainer, files)?; // 将分词器保存在一个包含所有配置和词汇的文件中 tokenizer.save("data/tokenizer-wiki.json", false)?; // 从该文件重新加载训练好的分词器 let mut tokenizer = Tokenizer::from_file("data/tokenizer-wiki.json")?; // 将训练好的分词器应用于任何想要的文本 // encode 方法会在文本上应用完整的分词器流程,返回一个 `Encoding` 对象 let output = tokenizer.encode("Hello, y'all! How are you 😁 ?", true)?; // Encoding 对象的 `tokens` 属性包含了文本的分词结果: println!("{:?}", output.get_tokens()); // ["Hello", ",", "y", "'", "all", "!", "How", "are", "you", "[UNK]", "?",] // Encoding 对象的 `id` 属性包含分词器词汇表中每个 token 的索引 println!("{:?}", output.get_ids()); // [27253, 16, 93, 11, 5097, 5, 7961, 5112, 6218, 0, 35] // Tokenizers库的一个重要特性是它具有完整的对齐跟踪功能 // `Encoding` 对象的 `offsets` 属性包含给定标记对应的原始句子部分 // 假设想找回导致 `"[UNK]"` 标记出现的原因,即列表中索引为 9 的标记, // 只需询问该索引处的偏移量即可。 println!("{:?}", output.get_offsets()[9]); // (26, 30) // 获取与原句中的表情符号对应的索引 let sentence = "Hello, y'all! How are you 😁 ?"; println!("{}", &sentence[26..30]); // "😁" // 后处理 // 当构建分词器时,将 `"[CLS]"` 和 `"[SEP]"` 设置在特殊标记列表的第1和第2位置 // 所以这应该是它们的 ID。 println!("{}", tokenizer.token_to_id("[SEP]").unwrap()); // 2 // 设置 BERT 后处理标记: let special_tokens = vec![ ("[CLS]", tokenizer.token_to_id("[CLS]").unwrap()), ("[SEP]", tokenizer.token_to_id("[SEP]").unwrap()), ]; // 指定了单个句子的模板:它们应该具有 `"[CLS] $A [SEP]"` 的形式 // 其中 `$A` 代表输入的句子 // 然后,为句子对指定模板,其形式应为 `"[CLS] $A [SEP] $B [SEP]"` // 其中 `$A` 代表第一句, `$B` 代表第二句 // 模板中添加的 `:1` 代表对输入的每个部分所需的 `type IDs` // 默认为 0 tokenizer.with_post_processor( TemplateProcessing::builder() .try_single("[CLS] $A [SEP]") .unwrap() .try_pair("[CLS] $A [SEP] $B:1 [SEP]:1") .unwrap() // 在分词器的词汇表中指定了使用的特殊标记及其ID .special_tokens(special_tokens) .build()?, ); // 开始使用分词器 let output = tokenizer.encode("Hello, y'all! How are you 😁 ?", true)?; println!("{:?}", output.get_tokens()); // ["[CLS]", "Hello", ",", "y", "'", "all", "!", "How", "are", "you", "[UNK]", "?", "[SEP]"] // 可以检查每个令牌所归属的类型ID是否正确 // 如果将分词器保存为 `Tokenizer.save` ,后处理器也将被保存 println!("{:?}", output.get_type_ids()); // [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1] // 批量处理文本 let output = tokenizer.encode_batch(vec!["Hello, y'all!", "How are you 😁 ?"], true)?; // 只要内存够用,可以同时处理任意数量的文本 let output = tokenizer.encode_batch( vec![ ("Hello, y'all!", "How are you 😁 ?"), ("Hello to you too!", "I'm fine, thank you!"), ], true, )?; // 当对多个句子进行编码时 // 可以使用 `Tokenizer.with_padding` 自动将输出填充到最长的句子 // 可以设置填充的位置(默认为右侧) tokenizer.with_padding(Some(PaddingParams { pad_id: 3, pad_token: "[PAD]".to_string(), ..PaddingParams::default() })); let output = tokenizer.encode_batch(vec!["Hello, y'all!", "How are you 😁 ?"], true)?; println!("{:?}", output[1].get_tokens()); // ["[CLS]", "How", "are", "you", "[UNK]", "?", "[SEP]", "[PAD]"] // 在这种情况下,分词器生成的 `attention mask` 会考虑到填充 println!("{:?}", output[1].get_attention_mask()); // [1, 1, 1, 1, 1, 1, 1, 0] } ``
本文从自然语言分词器的基本工作流出发,通过阅读 HuggingFace 的 tokenizers 库来进一步了解分词器的实现机制。本篇围绕整体架构,而非具体实现细节。下一篇将深入 BPE 分词模型细节中。
以上就是本篇的全部内容。感谢阅读。