refactor(tokenizer): 整理 tokenizer，添加新 bpe 实现

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

Cargo.lock

@@ -7,6 +7,5 @@ authors = ["YdrMaster <[email protected]>"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-common = { path = "../common" }
 memmap2.workspace = true
 patricia_tree = "0.8"

tokenizer/Cargo.toml

@@ -1,5 +1,4 @@
-use crate::{decode_with_ascii, Tokenizer};
-use common::utok;
+use crate::{decode_with_ascii, utok, Tokenizer};
 use std::{io::Result, path::Path};
 
 /// 由 tokenizer.model 文件定义的 bpe 分词器。
@@ -131,37 +130,37 @@ impl Tokenizer for BPE {
     }
 }
 
-#[test]
-fn read_tokenizer() {
-    let Some(model_dir) = common::test_model::find() else {
-        return;
-    };
-    println!("model_dir: {}", model_dir.display());
-
-    if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
-        for i in 0..bpe.offsets.len() {
-            println!("{}: {}", bpe.get_piece(i as utok), bpe.get_score(i as utok));
-        }
-    }
-}
-
-#[test]
-fn once_upon_a_time() {
-    let Some(model_dir) = common::test_model::find() else {
-        return;
-    };
-    println!("model_dir: {}", model_dir.display());
-
-    use std::time::Instant;
-    if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
-        const PROMPT: &str = "▁Once▁upon▁a▁time,";
-        let tokens = bpe.encode(PROMPT);
-        let t0 = Instant::now();
-        for _ in 0..10000 {
-            let _tokens = bpe.encode(PROMPT);
-        }
-        let t1 = Instant::now();
-        println!("{:?}", t1 - t0);
-        assert_eq!(tokens, &[9038, 2501, 263, 931, 29892]);
-    }
-}
+// #[test]
+// fn read_tokenizer() {
+//     let Some(model_dir) = common::test_model::find() else {
+//         return;
+//     };
+//     println!("model_dir: {}", model_dir.display());
+
+//     if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
+//         for i in 0..bpe.offsets.len() {
+//             println!("{}: {}", bpe.get_piece(i as utok), bpe.get_score(i as utok));
+//         }
+//     }
+// }
+
+// #[test]
+// fn once_upon_a_time() {
+//     let Some(model_dir) = common::test_model::find() else {
+//         return;
+//     };
+//     println!("model_dir: {}", model_dir.display());
+
+//     use std::time::Instant;
+//     if let Ok(bpe) = BPE::from_model_file(model_dir.join("tokenizer.model")) {
+//         const PROMPT: &str = "▁Once▁upon▁a▁time,";
+//         let tokens = bpe.encode(PROMPT);
+//         let t0 = Instant::now();
+//         for _ in 0..10000 {
+//             let _tokens = bpe.encode(PROMPT);
+//         }
+//         let t1 = Instant::now();
+//         println!("{:?}", t1 - t0);
+//         assert_eq!(tokens, &[9038, 2501, 263, 931, 29892]);
+//     }
+// }

tokenizer/src/bpe.rs

@@ -1,8 +1,11 @@
 mod bpe;
+mod new_bpe;
 mod normalizer;
 mod vocab_txt;
 
-use common::utok;
+/// `utok` for token id.
+#[allow(non_camel_case_types)]
+pub type utok = u32;
 
 pub trait Tokenizer {
     fn vocab_size(&self) -> usize;
@@ -14,7 +17,7 @@ pub use bpe::BPE;
 pub use normalizer::{BPECommonNormalizer, Normalizer};
 pub use vocab_txt::VocabTxt;
 
-const fn decode_with_ascii<'a>(piece: &'a str) -> &'a str {
+const fn decode_with_ascii(piece: &str) -> &str {
     // 预填充 ASCII 码表的所有字符
     const BYTES: [u8; 256] = {
         let mut ans = [0; 256];

tokenizer/src/lib.rs

@@ -0,0 +1,65 @@
+#![allow(unused)]
+
+use crate::utok;
+use std::{pin::Pin, ptr::NonNull};
+
+pub struct NewBpe {
+    vocabs: Pin<Box<[u8]>>,
+    token_to_piece: Box<[Token]>,
+    piece_to_token: Box<[utok]>,
+}
+
+struct Token {
+    ptr: NonNull<u8>,
+    len: u32,
+    score: f32,
+}
+
+impl AsRef<str> for Token {
+    #[inline]
+    fn as_ref(&self) -> &str {
+        use std::{slice::from_raw_parts, str::from_utf8_unchecked};
+        unsafe { from_utf8_unchecked(from_raw_parts(self.ptr.as_ptr(), self.len as _)) }
+    }
+}
+
+impl NewBpe {
+    pub fn new<'a>(
+        vocabs: impl IntoIterator<Item = &'a str>,
+        scores: impl Iterator<Item = f32>,
+        vocab_size_hint: usize,
+    ) -> Self {
+        let mut text_buf = Vec::with_capacity(vocab_size_hint * 4);
+        let mut token_to_piece = Vec::<(usize, usize)>::with_capacity(vocab_size_hint);
+
+        for vocab in vocabs.into_iter() {
+            let vocab = vocab.as_bytes();
+            let off = text_buf.len();
+            let len = vocab.len();
+            text_buf.extend_from_slice(vocab);
+            token_to_piece.push((off, len));
+        }
+        let vocab_size = token_to_piece.len();
+
+        let vocabs = unsafe { Pin::new_unchecked(text_buf.into_boxed_slice()) };
+        let token_to_piece = token_to_piece
+            .into_iter()
+            .zip(scores)
+            .map(|((off, len), score)| Token {
+                ptr: unsafe { NonNull::new_unchecked(vocabs.as_ptr().add(off).cast_mut()) },
+                len: len as _,
+                score,
+            })
+            .collect::<Box<[_]>>();
+        assert_eq!(token_to_piece.len(), vocab_size);
+
+        let mut piece_to_token = (0..vocab_size as utok).collect::<Box<[_]>>();
+        piece_to_token.sort_by_key(|&i| token_to_piece[i as usize].as_ref());
+
+        Self {
+            vocabs,
+            token_to_piece,
+            piece_to_token,
+        }
+    }
+}

tokenizer/src/new_bpe.rs

@@ -1,5 +1,4 @@
-use crate::{decode_with_ascii, Tokenizer};
-use common::utok;
+use crate::{decode_with_ascii, utok, Tokenizer};
 use memmap2::Mmap;
 use patricia_tree::PatriciaMap;
 use std::{fs::File, io::Result, path::Path};