feat(memory): 千问字节级 BPE 分词器 + 与上游逐条对齐的回归测试

为「千问嵌入模型导出 ONNX 并内嵌」的 Go 侧准备。CLIP 那套 tokenizer 不能复用:
CLIP 是「小写化 + 空白规整 + 词表 BPE」,千问是 **GPT-2 式字节级 BPE**
(先按字节映射到安全 unicode,再对映射结果做合并),中文与空白输入的切分
完全不同。

实现中撞到两个与上游对齐的坑,都由测试暴露:

1. **`\s+(?!\S)` 的语义依赖正则回溯**,不是「等价于 `\s+`」。
   `\s+` 先贪婪吃完整段空白,发现后面是非空白导致 `(?!\S)` 失败,于是回退
   一个字符,**正好留下末尾一个空白**给前面以 ` ?` / `[^…]?` 开头的分支合并。
   这直接决定切分点:`"   leading"` 会切成 `"  "` + `" leading"`,
   而不是 `"   "` + `"leading"`。RE2 不支持 lookaround,近似改写必然对不上,
   所以改成按分支顺序**显式实现**(含 `\s*[\r\n]+` 的回溯语义)。
   实测:近似改写时 26 条里错 3 条,全部是空白串用例。

2. **Go 的 `\s` 只有 ASCII,且 regexp 不支持二进制属性 `\p{White_Space}`**
   (只支持 script/category,直接写会报 invalid character class range)。
   而上游 Rust regex 的 `\s` 正是 White_Space。改用 `unicode.IsSpace` 作为
   唯一判据,避免全角空格/NBSP/行分隔符的切分点漂移。

另外特殊 token(`<|im_start|>` 等 24 个 AddedToken)必须**整体优先匹配**并
按长度降序,否则会被 BPE 拆成子 token,模型收到的输入就变了——且不会报任何错。

验证:`testdata/qwen_tokenizer_reference.json` 由 HuggingFace 真实 tokenizer
生成(26 条用例,覆盖中/英/中英混排/数字/各类空白形态/标点/emoji/特殊 token/
长文本/空串/单字符边界),Go 实现逐条精确对齐。字节↔unicode 映射另测双射性
(有碰撞会让不同字节编成同一 token,静默产生错误输入)。
This commit is contained in:
JianFeeeee
2026-09-11 00:11:58 +08:00
parent b393b7072c
commit e985151da9
3 changed files with 2000 additions and 0 deletions

View File

@ -0,0 +1,161 @@
package qwen
import (
"encoding/json"
"os"
"path/filepath"
"testing"
)
// modelDir 是本地千问模型目录。不存在则跳过——参考数据已固化在 testdata,
// 但分词器本身要从 tokenizer.json 加载词表与 merges(11MB,不入库)。
const modelDir = "/home/newqqagent/models/models/qwen--Qwen3-VL-Embedding-2B/snapshots/master"
type tokenizerRef struct {
VocabSize int `json:"vocab_size"`
Cases []struct {
Text string `json:"text"`
IDs []int `json:"ids"`
Tokens []string `json:"tokens"`
} `json:"cases"`
AddedTokens []struct {
Content string `json:"content"`
ID int `json:"id"`
Special bool `json:"special"`
} `json:"added_tokens"`
}
func loadRef(t *testing.T) *tokenizerRef {
t.Helper()
raw, err := os.ReadFile(filepath.Join("testdata", "qwen_tokenizer_reference.json"))
if err != nil {
t.Fatalf("读取参考数据: %v", err)
}
var ref tokenizerRef
if err := json.Unmarshal(raw, &ref); err != nil {
t.Fatalf("解析参考数据: %v", err)
}
return &ref
}
func loadTokenizer(t *testing.T) *Tokenizer {
t.Helper()
if _, err := os.Stat(filepath.Join(modelDir, "tokenizer.json")); err != nil {
t.Skipf("模型目录不可用,跳过: %v", err)
}
tok, err := LoadTokenizer(modelDir)
if err != nil {
t.Fatalf("LoadTokenizer: %v", err)
}
return tok
}
// 与 HuggingFace 的真实 tokenizer 逐条对齐。
//
// 这是本包唯一的正确性判据:字节级 BPE 的失败模式是「看起来能跑但 token 不同」,
// 而 token 不同会让模型收到完全不同的输入,嵌入自然也就错了——不会报任何错。
// 所以必须拿真实输出对照,不能靠读代码断言。
func TestTokenizerMatchesReference(t *testing.T) {
ref := loadRef(t)
tok := loadTokenizer(t)
if got := tok.VocabSize(); got != ref.VocabSize {
t.Errorf("词表大小 = %d,参考 %d", got, ref.VocabSize)
}
failed := 0
for _, c := range ref.Cases {
got := tok.Encode(c.Text)
if !sameIDs(got, c.IDs) {
failed++
t.Errorf("不一致 text=%q\n got %v\n want %v", c.Text, got, c.IDs)
}
}
if failed > 0 {
t.Fatalf("%d/%d 条用例不一致", failed, len(ref.Cases))
}
}
func sameIDs(a, b []int) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
// 特殊 token 必须整体匹配:走 BPE 会被拆成子 token,模型看到的输入就变了。
func TestSpecialTokensMatchWhole(t *testing.T) {
ref := loadRef(t)
tok := loadTokenizer(t)
for _, at := range ref.AddedTokens {
if !at.Special {
continue
}
got, ok := tok.SpecialID(at.Content)
if !ok {
t.Errorf("特殊 token %q 未从 tokenizer.json 载入", at.Content)
continue
}
if got != at.ID {
t.Errorf("特殊 token %q id=%d,参考 %d", at.Content, got, at.ID)
}
// 单独出现时必须编码成恰好一个 id。
ids := tok.Encode(at.Content)
if len(ids) != 1 || ids[0] != at.ID {
t.Errorf("特殊 token %q 应整体编码为 [%d],实际 %v", at.Content, at.ID, ids)
}
}
}
// 最长优先:`<|im_start|>` 不能被更短的 `<|im_end|>` 之类前缀抢走。
func TestSpecialTokenLongestFirst(t *testing.T) {
tok := loadTokenizer(t)
text := "<|im_start|>user\n你好<|im_end|>"
ids := tok.Encode(text)
startID, _ := tok.SpecialID("<|im_start|>")
endID, _ := tok.SpecialID("<|im_end|>")
if len(ids) == 0 || ids[0] != startID {
t.Fatalf("应以 <|im_start|>(%d) 开头,实际 %v", startID, ids)
}
if last := ids[len(ids)-1]; last != endID {
t.Fatalf("应以 <|im_end|>(%d) 结尾,实际 %v", endID, ids)
}
}
// 空串与单字符边界。
func TestTokenizerEdgeCases(t *testing.T) {
tok := loadTokenizer(t)
if got := tok.Encode(""); len(got) != 0 {
t.Errorf("空串应产出 0 个 token,实际 %v", got)
}
for _, s := range []string{"a", "中", "1", " "} {
if got := tok.Encode(s); len(got) == 0 {
t.Errorf("%q 应至少产出 1 个 token", s)
}
}
}
// byteEnc 必须是双射:256 个字节映射到 256 个互不相同的码点。
// 有碰撞就会让不同字节编成同一个 token,静默产生错误输入。
func TestBytesToUnicodeBijective(t *testing.T) {
m := bytesToUnicode()
if len(m) != 256 {
t.Fatalf("映射应覆盖 256 个字节,实际 %d", len(m))
}
seen := map[rune]byte{}
for b, r := range m {
if prev, dup := seen[r]; dup {
t.Fatalf("码点冲突:字节 %d 与 %d 都映射到 %q", prev, b, r)
}
seen[r] = b
}
}