【Rust 研学】 | LLM 入门之旅 2 : BPE 算法

/// Base trait for Tokenizers to implement.
pub trait Tokenizer {
 // 提供对特殊标记的访问，这些特殊标记通常用于处理如句子开始、结束等特定功能。
    fn special_tokens(&self) -> &IndexMap<String, Token>;

 // 提供对合并规则的访问，这些规则定义了在训练过程中哪些标记被合并。
    fn merges(&self) -> &IndexMap<(Token, Token), Token>;

 // 提供对词汇表的访问，这是标记化过程中的关键数据结构。
    fn vocab(&self) -> &IndexMap<Token, Vec<u8>>;

 // 将文本转换成一系列标记ID。这是文本处理中的基础步骤，用于后续的处理如模型训练或文本生成。
    /// A Tokenizer can encode a string into a list of integers.
    fn encode(&self, text: &str) -> Vec<Token>;

 // 将标记ID序列转换回原始文本。这通常用于生成文本后的输出阶段，验证标记化过程的准确性或用户界面展示
    /// A Tokenizer can decode a list of integers into a string.
    fn decode(&self, ids: &[Token]) -> String;
}

/// A Tokenizer that can be trained.
pub trait Trainable: Tokenizer {
 // 允许从给定的文本中训练词汇表和合并规则，
 // `vocab_size` 指定了目标词汇表大小，
 // `verbose` 标志控制是否输出详细的训练信息。
    /// Train a vocabulary of size `vocab_size` in distinct Tokens from `text`.
    fn train(&mut self, text: &str, vocab_size: Token, verbose: bool);
}


pub trait Saveable: Tokenizer {
    fn pattern(&self) -> &str;
    /// Saves the tokenizer's model and vocabulary to two files
    fn save(&self, dir: &Path, prefix: &str) {
        // let dir = dir.as_ref();

        // Write the model file (used for loading the tokenizer later)
        // 省略实现
    }
}


pub trait Loadable: Tokenizer {
    fn set_pattern(&mut self, pattern: &str);
    fn set_special_tokens(&mut self, special_tokens: IndexMap<String, Token>);
    fn set_merges(&mut self, merges: IndexMap<(Token, Token), Token>);
    fn set_vocab(&mut self, vocab: IndexMap<Token, Vec<u8>>);

    /// Loads the tokenizer's model from a file.
    fn load(&mut self, model_file: &Path) {
  // 省略实现
 }
}

pub struct BasicTokenizer {
    special_tokens: IndexMap<String, Token>,
    merges: IndexMap<(Token, Token), Token>,
    vocab: IndexMap<Token, Vec<u8>>,
}


impl Tokenizer for BasicTokenizer {
    fn special_tokens(&self) -> &IndexMap<String, Token> {
        &self.special_tokens
    }

    fn merges(&self) -> &IndexMap<(Token, Token), Token> {
        &self.merges
    }

    fn vocab(&self) -> &IndexMap<Token, Vec<u8>> {
        &self.vocab
    }

    fn decode(&self, ids: &[Token]) -> String {
        // 将输入的标记ID序列转换成字符串
        // 通过遍历每个标记ID，从 `vocab` 映射中查找对应的字节序列
        // 然后将这些序列合并成一个完整的 UTF-8 字符串
        let text_bytes: Vec<u8> = ids
            .iter()
            .flat_map(|&idx| self.vocab[&idx].clone())
            .collect();
        // 将字节向量转换为字符串，
        // 这个方法会用特殊字符替换任何无效的 UTF-8 序列。
        String::from_utf8_lossy(&text_bytes).into_owned()
    }

    fn encode(&self, text: &str) -> Vec<Token> {
        // 将输入文本转换为其字节表示形式的序列
        let text_bytes = text.as_bytes();
        let mut ids: Vec<Token> = text_bytes.iter().map(|&b| b as Token).collect();
        while ids.len() >= 2 {
            // 找出序列中最常见的相邻标记对
            // 根据 `merges` 选择最低的合并索引来合并标记
            let stats = get_stats(&ids);

            let pair_opt = stats
                .keys()
                .filter_map(|&pair| self.merges.get(&pair).map(|_| pair))
                .min_by_key(|&pair| self.merges[&pair]);

   // 循环进行直到没有可合并的标记对为止
            match pair_opt {
                None => break, // If there are no more merges available, break
                Some(pair) => {
                    // Otherwise, merge the best pair (lowest merge index)
                    let idx = self.merges[&pair];
                    ids = merge(&ids, pair, idx);
                }
            };
        }
        ids
    }
}

impl Trainable for BasicTokenizer {
    /// 根据提供的文本训练一个大小为 `vocab_size` 的词汇表。
    /// 该方法实现了字节对编码（Byte Pair Encoding, BPE）算法，迭代地找出并合并最频繁的相邻标记对，
    /// 直到词汇表达到所需的大小。
    fn train(&mut self, text: &str, vocab_size: Token, verbose: bool) {
        // 确保请求的词汇表大小至少为256，以容纳所有单字节字符。
        assert!(vocab_size >= 256, "词汇表大小必须至少为256");

        // 计算需要创建的新标记数量，词汇表大小减去256（基础单字节字符的数量）。
        let num_merges = vocab_size - 256;

        // 将输入文本预处理为字节序列，每个字节视为一个初始标记。
        let text_bytes = text.as_bytes();
        let mut ids: Vec<Token> = text_bytes.iter().map(|&b| b as Token).collect();

        // 初始化合并记录和词汇表。
        let mut merges: IndexMap<(Token, Token), Token> = IndexMap::new();
        let mut vocab: IndexMap<Token, Vec<u8>> = (0..256).map(|idx| (idx, vec![idx as u8])).collect();

        // 迭代合并最常见的标记对，直到达到预设的词汇表大小或没有可合并的标记对。
        for i in 0..num_merges {
            // 计算当前所有相邻标记对的出现频率。
            let stats = get_stats(&ids);
            // 找到出现次数最多的标记对。
            let pair = get_max_entry(&stats).unwrap().0;

            // 为合并后的新标记分配一个新的ID（从256开始）。
            let idx = 256 + i;
            // 在标记序列中替换所有出现的该标记对为新标记。
            ids = merge(&ids, *pair, idx);
            
            // 保存合并规则和更新词汇表。
            merges.insert(*pair, idx);
            vocab.insert(
                idx,
                [vocab[&pair.0].clone(), vocab[&pair.1].clone()].concat(),
            );

            // 如果设置为详细模式，则打印每次合并的详细信息。
            if verbose {
                println!(
                    "merge {}/{}: {:?} -> {} ({:?}) had {} occurrences",
                    i + 1,
                    num_merges,
                    pair,
                    idx,
                    vocab[&idx],
                    stats[pair]
                );
            }
        }

        // 保存实例变量。
        self.merges = merges;
        self.vocab = vocab;
    }
}

pub const GPT2_SPLIT_PATTERN: &str = r"'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+";
pub const GPT4_SPLIT_PATTERN: &str = r"'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+";

pub enum AllowedSpecial {
    All, // 允许在编码中使用所有特殊标记
    None, // 忽略所有特殊标记，将其视为普通文本进行编码
    NoneRaise, // 如果在编码过程中遇到特殊标记则引发错误
    Set(HashSet<String>), // 仅允许指定的特殊标记集合
}

pub trait RegexTokenizerTrait: Tokenizer {
 // 获取编译后的正则表达式对象
    fn compiled_pattern(&self) -> &Regex; 
 // 获取反向特殊标记映射
    fn inverse_special_tokens(&self) -> &IndexMap<Token, String>; 
 // 将文本编码成标记ID序列
    fn encode(&self, text: &str) -> Vec<Token> { 
     // 如果遇到特殊标记默认情况下引发错误
        self.encode_special(text, AllowedSpecial::NoneRaise) 
    }
 // 将标记ID序列解码成字符串
    fn decode(&self, ids: &[Token]) -> String { 
        // 默认实现，解码逻辑，包括处理特殊标记
    }
 // 编码任意忽略的特殊标记
 fn encode_ordinary(&self, text: &str) -> Vec<Token> {
   // 默认实现
  }
 // 根据 allowed_special 参数处理特殊标记的编码逻辑
    fn encode_special(&self, text: &str, allowed_special: AllowedSpecial) -> Vec<Token> {
       // 默认实现 
    }
}

pub struct RegexTokenizerStruct {
 // 使用的正则表达式模式字符串
    pattern: String, 
    // 编译后的正则表达式对象
    compiled_pattern: Regex, 
    // 特殊标记映射
    special_tokens: IndexMap<String, Token>,
    // 反向特殊标记映射 
    inverse_special_tokens: IndexMap<Token, String>, 
    // 合并规则映射
    merges: IndexMap<(Token, Token), Token>, 
    // 词汇表
    vocab: IndexMap<Token, Vec<u8>>, 
}
impl Tokenizer for RegexTokenizerStruct {
 // 实现
}
impl Trainable for RegexTokenizerStruct {
 // 实现
}