mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-04 23:00:30 +08:00
fix(tokenizer): load cl100k BPE table from disk instead of failing silently offline (#17712)
## Summary RAGFlow's Go tokenizer silently returned **0 tokens for every string** whenever the `cl100k_base` BPE table could not be loaded — which is the normal case for an offline/air-gapped Go server. This PR makes the loader resolve the table from disk (where RAGFlow actually ships it) and fail loudly when it is genuinely missing. ## Root cause `tiktoken-go`'s stock loader downloads the encoding table over HTTP and caches it under `TIKTOKEN_CACHE_DIR`. That does not work for RAGFlow: - `TIKTOKEN_CACHE_DIR` is exported **only inside the Python process** (`common/token_utils.py`). `docker/entrypoint.sh` launches the Go binary (`bin/ragflow_server`) from a shell, so the Go process never inherits the variable. - The Dockerfile *does* ship the table (under its sha1 name in the working directory), but nothing told the Go side to look there. - Reaching `openaipublic.blob.core.windows.net` at runtime is not an option for air-gapped installs, and is unreliable where that host is blocked. The failure was **silent**: `NumTokensFromString` returns `0` when the encoder fails to build, and a `sync.Once` memoizes that error for the process lifetime. Every token count became `0`, so chunk merging never crossed its token budget and an entire document collapsed into a single chunk. Python has no such failure mode because its encoder is built at import time (a missing table aborts startup instead of degrading). ## Fix Register a local-only `BpeLoader` via `tiktoken.SetBpeLoader` (`internal/tokenizer/bpe_loader.go`) that resolves the table from disk **only**, in priority order: 1. `TIKTOKEN_CACHE_DIR` / `DATA_GYM_CACHE_DIR` (honored so operators who already configured one keep working). 2. The working directory, the executable's directory, and all of their ancestors — matching the Dockerfile layout (table under its sha1 name in the install root). 3. A `ragflow_deps/<basename>` checkout produced by `ragflow_deps/download_deps.py`. It **never performs network I/O**. When nothing is found it returns an error listing every path it tried (pointing at `download_deps.py` or `TIKTOKEN_CACHE_DIR`), so a genuinely missing table fails loudly instead of degrading to zero. ## Test plan - `internal/tokenizer/bpe_loader_test.go` (unit tier, runs under `bash build.sh --test ./internal/tokenizer/...`): - Loader reads from `TIKTOKEN_CACHE_DIR`, `DATA_GYM_CACHE_DIR`, the sha1-named file in the working dir, and the bundled `ragflow_deps/` name. - Explicit cache dir wins over the bundled vocab. - A malformed table is reported as an error rather than skipped. - A genuinely missing table reports the candidates it tried (no network attempt). - `NumTokensFromString` matches Python-derived anchors (`""`→0, `"hello"`→1, `"hello world"`→2, `"hello, world!"`→4, `"世界"`→3, `"Hello 世界 🌍"`→8, `"RAGFlow"`→3). ## Notes - `.github/workflows/tests.yml` currently excludes `internal/tokenizer` from `go test`, so these tests do not run in CI. The tokenizer fix is exercised in CI indirectly via the chunker package once a token-count-sensitive parity case lands (tracked separately). Consider including `internal/tokenizer` in CI as a follow-up. - Supported deployments already ship the table (`download_deps.py` → `ragflow_deps/cl100k_base.tiktoken`; Dockerfile → `<sha1>` in cwd), so no `ENV` change is required for the fix to take effect. Setting `ENV TIKTOKEN_CACHE_DIR` in the Dockerfile remains a cheap belt-and-suspenders hardening that can be done separately. 🤖 Generated with [CodeBuddy Code](https://cnb.cool/codebuddy) --------- Co-authored-by: CodeBuddy <noreply@codebuddy.ai> Co-authored-by: CodeBuddy Code <noreply@cnb.cool> Co-authored-by: CodeBuddy <noreply@tencent.com>
This commit is contained in:
216
internal/tokenizer/bpe_loader.go
Normal file
216
internal/tokenizer/bpe_loader.go
Normal file
@@ -0,0 +1,216 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tokenizer
|
||||
|
||||
// Offline BPE table loading for tiktoken.
|
||||
//
|
||||
// RAGFlow ships the cl100k_base table on disk (Dockerfile drops it into the
|
||||
// working directory under its sha1 name; download_deps.py writes it to
|
||||
// ragflow_deps/). tiktoken-go's stock loader instead downloads it over HTTP and
|
||||
// relies on TIKTOKEN_CACHE_DIR, which the Go server never inherits, so a
|
||||
// missing table degrades every token count to 0. This loader resolves the
|
||||
// table from disk only: it performs no network I/O, and when nothing is found
|
||||
// it reports every path it tried.
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"os"
|
||||
"path"
|
||||
"path/filepath"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"ragflow/internal/common"
|
||||
|
||||
"github.com/pkoukk/tiktoken-go"
|
||||
)
|
||||
|
||||
func init() {
|
||||
tiktoken.SetBpeLoader(localBpeLoader{})
|
||||
}
|
||||
|
||||
// expectedBpeHashes maps a tiktoken table URL to the SHA-1 of its canonical
|
||||
// on-disk content. We only ship cl100k_base today; entries here let the loader
|
||||
// reject a corrupt or tampered file instead of trusting it. Unknown URLs are
|
||||
// loaded without a digest check (defense-in-depth, not a hard gate).
|
||||
//
|
||||
// NOTE: this is the digest of the file *contents*, not the tiktoken cache
|
||||
// filename. tiktoken-go names its cached file by sha1(bpeURL)
|
||||
// (223921b76ee99bde995b7ff738513eef100fb51d18c93597a113bcffe865b2a7 for
|
||||
// cl100k_base); that value identifies the path, while the value below verifies
|
||||
// the bytes we actually load. Compute it from the table shipped by
|
||||
// ragflow_deps/download_deps.py: `sha1sum cl100k_base.tiktoken`.
|
||||
var expectedBpeHashes = map[string]string{
|
||||
"https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken": "6494e42d5aad2bbb441ea9793af9e7db335c8d9c",
|
||||
}
|
||||
|
||||
// localBpeLoader resolves tiktoken BPE tables from the local filesystem.
|
||||
type localBpeLoader struct{}
|
||||
|
||||
// LoadTiktokenBpe implements tiktoken.BpeLoader.
|
||||
//
|
||||
// bpeURL is the upstream table URL that tiktoken-go would otherwise download;
|
||||
// here it serves only to derive the file names to look for.
|
||||
func (localBpeLoader) LoadTiktokenBpe(bpeURL string) (map[string]int, error) {
|
||||
candidates := bpeCandidatePaths(bpeURL)
|
||||
for _, candidate := range candidates {
|
||||
contents, err := os.ReadFile(candidate)
|
||||
if err != nil {
|
||||
// Only a missing candidate is skippable; a permission or I/O
|
||||
// failure must not be masked as "not found".
|
||||
if os.IsNotExist(err) {
|
||||
continue
|
||||
}
|
||||
return nil, fmt.Errorf("reading BPE table %s: %w", candidate, err)
|
||||
}
|
||||
// Integrity check: for tables we ship, a digest mismatch means the
|
||||
// file is corrupt or tampered with. Refuse to load it rather than
|
||||
// skipping to the next candidate — a different candidate holds the
|
||||
// same (wrong) content, and masking the failure would defeat the
|
||||
// check. This mirrors the malformed-table path just below.
|
||||
if want, ok := expectedBpeHashes[bpeURL]; ok {
|
||||
if got := fmt.Sprintf("%x", sha1.Sum(contents)); got != want {
|
||||
return nil, fmt.Errorf("BPE table %s digest mismatch (got %s, want %s); refusing to load a corrupt or tampered file", candidate, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
ranks, err := parseBpeTable(contents)
|
||||
if err != nil {
|
||||
// A file that exists but does not parse is a corrupt download or a
|
||||
// name collision. Continuing to the next candidate would mask it.
|
||||
return nil, fmt.Errorf("BPE table %s is malformed: %w", candidate, err)
|
||||
}
|
||||
return ranks, nil
|
||||
}
|
||||
|
||||
err := fmt.Errorf(
|
||||
"no local BPE table for %s; run `uv run ragflow_deps/download_deps.py` or set TIKTOKEN_CACHE_DIR to the directory holding the table; tried: %s",
|
||||
bpeURL, strings.Join(candidates, ", "))
|
||||
// Logged as well as returned: tiktoken-go propagates this to GetEncoding,
|
||||
// whose error NumTokensFromString discards to keep returning 0.
|
||||
common.Error("cl100k BPE table not found; every token count will be 0", err)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// bpeCandidatePaths lists, in priority order, every local path that may hold
|
||||
// the table for bpeURL.
|
||||
//
|
||||
// Explicit configuration wins, then the directories RAGFlow actually ships the
|
||||
// table in. Both the working directory and the executable's directory are
|
||||
// walked upwards: the server runs with the working directory set to the
|
||||
// installation root, while `go test` runs from a package subdirectory.
|
||||
func bpeCandidatePaths(bpeURL string) []string {
|
||||
cacheName := fmt.Sprintf("%x", sha1.Sum([]byte(bpeURL)))
|
||||
// download_deps.py stores the table under the URL's own basename.
|
||||
bundledName := path.Base(bpeURL)
|
||||
|
||||
var paths []string
|
||||
seen := make(map[string]struct{})
|
||||
add := func(p string) {
|
||||
if _, dup := seen[p]; dup {
|
||||
return
|
||||
}
|
||||
seen[p] = struct{}{}
|
||||
paths = append(paths, p)
|
||||
}
|
||||
|
||||
// Honour both variables tiktoken-go itself reads, so an operator who has
|
||||
// already configured one keeps working.
|
||||
for _, env := range []string{"TIKTOKEN_CACHE_DIR", "DATA_GYM_CACHE_DIR"} {
|
||||
if dir := strings.TrimSpace(os.Getenv(env)); dir != "" {
|
||||
add(filepath.Join(dir, cacheName))
|
||||
}
|
||||
}
|
||||
|
||||
for _, root := range searchRoots() {
|
||||
// Same layout the Dockerfile creates: the table sits in the
|
||||
// installation root under its sha1 name.
|
||||
add(filepath.Join(root, cacheName))
|
||||
// download_deps.py writes the table into ragflow_deps/ under its
|
||||
// download name; a developer checkout that has run it but never
|
||||
// started the Python side only has this copy.
|
||||
add(filepath.Join(root, "ragflow_deps", bundledName))
|
||||
}
|
||||
|
||||
return paths
|
||||
}
|
||||
|
||||
// searchRoots returns the working directory and the executable's directory
|
||||
// together with all of their ancestors.
|
||||
func searchRoots() []string {
|
||||
var roots []string
|
||||
seen := make(map[string]struct{})
|
||||
for _, start := range startingDirs() {
|
||||
for dir := start; ; {
|
||||
if _, dup := seen[dir]; !dup {
|
||||
seen[dir] = struct{}{}
|
||||
roots = append(roots, dir)
|
||||
}
|
||||
parent := filepath.Dir(dir)
|
||||
if parent == dir {
|
||||
break
|
||||
}
|
||||
dir = parent
|
||||
}
|
||||
}
|
||||
return roots
|
||||
}
|
||||
|
||||
func startingDirs() []string {
|
||||
var dirs []string
|
||||
if wd, err := os.Getwd(); err == nil {
|
||||
dirs = append(dirs, wd)
|
||||
}
|
||||
if exe, err := os.Executable(); err == nil {
|
||||
if resolved, err := filepath.EvalSymlinks(exe); err == nil {
|
||||
exe = resolved
|
||||
}
|
||||
dirs = append(dirs, filepath.Dir(exe))
|
||||
}
|
||||
return dirs
|
||||
}
|
||||
|
||||
// parseBpeTable decodes tiktoken's on-disk format: one
|
||||
// "<base64 token> <rank>" pair per line.
|
||||
func parseBpeTable(contents []byte) (map[string]int, error) {
|
||||
ranks := make(map[string]int)
|
||||
for i, line := range strings.Split(string(contents), "\n") {
|
||||
line = strings.TrimRight(line, "\r")
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
token, rank, ok := strings.Cut(line, " ")
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("line %d: expected \"<token> <rank>\"", i+1)
|
||||
}
|
||||
decoded, err := base64.StdEncoding.DecodeString(token)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("line %d: %w", i+1, err)
|
||||
}
|
||||
value, err := strconv.Atoi(rank)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("line %d: %w", i+1, err)
|
||||
}
|
||||
ranks[string(decoded)] = value
|
||||
}
|
||||
if len(ranks) == 0 {
|
||||
return nil, fmt.Errorf("table is empty")
|
||||
}
|
||||
return ranks, nil
|
||||
}
|
||||
77
internal/tokenizer/bpe_loader_anchors_test.go
Normal file
77
internal/tokenizer/bpe_loader_anchors_test.go
Normal file
@@ -0,0 +1,77 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
//go:build manual
|
||||
|
||||
package tokenizer
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestNumTokensFromString_MatchesPythonAnchors pins exact counts taken from the
|
||||
// Python reference suite (test/unit_test/common/test_token_utils.py:28-49) and
|
||||
// from common.token_utils.num_tokens_from_string for the CJK cases. The corpus
|
||||
// was later expanded to ~24 entries spanning ASCII, punctuation, digits,
|
||||
// whitespace, newlines, CJK, emoji, mixed-language, and code-like strings, all
|
||||
// recomputed against tiktoken's cl100k encoder.
|
||||
//
|
||||
// Exact values matter more than they look. NumTokensFromString swallows loader
|
||||
// errors and returns 0, so an assertion of the form "> 0" passes for an empty
|
||||
// string and fails to notice a dead encoder — which is precisely how the
|
||||
// offline breakage stayed invisible. Pinning the numbers also catches loading a
|
||||
// structurally valid but wrong table.
|
||||
//
|
||||
// This test needs the real cl100k_base table on disk (TIKTOKEN_CACHE_DIR,
|
||||
// the Dockerfile's /ragflow/<sha1> file, or ragflow_deps/cl100k_base.tiktoken),
|
||||
// so it is tagged `manual` and runs only under `build.sh --test-manual`,
|
||||
// which the docker builder provisions with /usr/share/infinity/resource.
|
||||
func TestNumTokensFromString_MatchesPythonAnchors(t *testing.T) {
|
||||
anchors := []struct {
|
||||
in string
|
||||
want int
|
||||
}{
|
||||
{"", 0},
|
||||
{"hello", 1},
|
||||
{"hello world", 2},
|
||||
{"hello, world!", 4},
|
||||
{"世界", 3},
|
||||
{"Hello 世界 🌍", 8},
|
||||
{"RAGFlow", 3},
|
||||
{"1234567890", 4},
|
||||
{"a b", 3},
|
||||
{"hello\nworld", 3},
|
||||
{"user@example.com", 3},
|
||||
{"https://example.com/path?x=1", 9},
|
||||
{"func main() {}", 4},
|
||||
{"aaaaaaaaaa", 2},
|
||||
{"中文字符测试", 4},
|
||||
{"🚀🔥", 6},
|
||||
{"state-of-the-art", 4},
|
||||
{`"quoted"`, 3},
|
||||
{"The quick brown fox jumps over the lazy dog.", 10},
|
||||
{"Café naïve résumé", 8},
|
||||
{"x² + y² = z²", 8},
|
||||
{"混合 English 和 中文 的 sentence。", 10},
|
||||
{"tokenization is the process of splitting text into tokens", 10},
|
||||
{"人工智能正在改变世界,这是毫无疑问的事实。", 25},
|
||||
{"SELECT * FROM users WHERE id = 1;", 10},
|
||||
{"こんにちは世界", 4},
|
||||
}
|
||||
for _, tc := range anchors {
|
||||
if got := NumTokensFromString(tc.in); got != tc.want {
|
||||
t.Errorf("NumTokensFromString(%q) = %d, want %d", tc.in, got, tc.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
201
internal/tokenizer/bpe_loader_test.go
Normal file
201
internal/tokenizer/bpe_loader_test.go
Normal file
@@ -0,0 +1,201 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"crypto/sha1"
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
const testBpeURL = "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken"
|
||||
|
||||
// writeBpeTable writes a minimal well-formed BPE table. The rank values are
|
||||
// arbitrary markers so a test can tell which file the loader actually read.
|
||||
func writeBpeTable(t *testing.T, path string, marker int) {
|
||||
t.Helper()
|
||||
if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
|
||||
t.Fatalf("mkdir %s: %v", filepath.Dir(path), err)
|
||||
}
|
||||
line := fmt.Sprintf("%s %d\n", base64.StdEncoding.EncodeToString([]byte("hello")), marker)
|
||||
if err := os.WriteFile(path, []byte(line), 0o644); err != nil {
|
||||
t.Fatalf("write %s: %v", path, err)
|
||||
}
|
||||
// Register this synthetic table's digest so the SHA-1 integrity gate in
|
||||
// bpe_loader.go accepts it. This also exercises the verification path
|
||||
// instead of disabling it.
|
||||
expectedBpeHashes[testBpeURL] = fmt.Sprintf("%x", sha1.Sum([]byte(line)))
|
||||
}
|
||||
|
||||
func cacheFileName(url string) string {
|
||||
return fmt.Sprintf("%x", sha1.Sum([]byte(url)))
|
||||
}
|
||||
|
||||
// isolate moves the process into an empty directory and clears both cache
|
||||
// environment variables, so a test sees only the files it creates itself.
|
||||
// Without this the real repository checkout — which does ship the table — would
|
||||
// satisfy every lookup and hide ordering bugs.
|
||||
func isolate(t *testing.T) string {
|
||||
t.Helper()
|
||||
dir := t.TempDir()
|
||||
t.Setenv("TIKTOKEN_CACHE_DIR", "")
|
||||
t.Setenv("DATA_GYM_CACHE_DIR", "")
|
||||
t.Chdir(dir)
|
||||
return dir
|
||||
}
|
||||
|
||||
func TestLocalBpeLoader_ReadsTiktokenCacheDir(t *testing.T) {
|
||||
isolate(t)
|
||||
cache := t.TempDir()
|
||||
writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 7)
|
||||
t.Setenv("TIKTOKEN_CACHE_DIR", cache)
|
||||
|
||||
ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got := ranks["hello"]; got != 7 {
|
||||
t.Errorf("rank from TIKTOKEN_CACHE_DIR = %d, want 7", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalBpeLoader_ReadsDataGymCacheDir(t *testing.T) {
|
||||
isolate(t)
|
||||
cache := t.TempDir()
|
||||
writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 8)
|
||||
t.Setenv("DATA_GYM_CACHE_DIR", cache)
|
||||
|
||||
ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got := ranks["hello"]; got != 8 {
|
||||
t.Errorf("rank from DATA_GYM_CACHE_DIR = %d, want 8", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The production image has no TIKTOKEN_CACHE_DIR set: Dockerfile drops the
|
||||
// table straight into the working directory under its sha1 name, and
|
||||
// entrypoint.sh starts the Go binary from a shell, so nothing exports the
|
||||
// variable that common/token_utils.py sets inside the Python process.
|
||||
func TestLocalBpeLoader_ReadsSha1FileFromWorkingDirectory(t *testing.T) {
|
||||
dir := isolate(t)
|
||||
writeBpeTable(t, filepath.Join(dir, cacheFileName(testBpeURL)), 9)
|
||||
|
||||
ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got := ranks["hello"]; got != 9 {
|
||||
t.Errorf("rank from working directory = %d, want 9", got)
|
||||
}
|
||||
}
|
||||
|
||||
// A developer checkout that has run download_deps.py but never started the
|
||||
// Python side only has the file under its download name.
|
||||
func TestLocalBpeLoader_ReadsBundledVocabFromAncestor(t *testing.T) {
|
||||
dir := isolate(t)
|
||||
writeBpeTable(t, filepath.Join(dir, "ragflow_deps", "cl100k_base.tiktoken"), 10)
|
||||
nested := filepath.Join(dir, "internal", "tokenizer")
|
||||
if err := os.MkdirAll(nested, 0o755); err != nil {
|
||||
t.Fatalf("mkdir: %v", err)
|
||||
}
|
||||
t.Chdir(nested)
|
||||
|
||||
ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got := ranks["hello"]; got != 10 {
|
||||
t.Errorf("rank from bundled vocab = %d, want 10", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalBpeLoader_CacheDirWinsOverBundledVocab(t *testing.T) {
|
||||
dir := isolate(t)
|
||||
writeBpeTable(t, filepath.Join(dir, "ragflow_deps", "cl100k_base.tiktoken"), 10)
|
||||
cache := t.TempDir()
|
||||
writeBpeTable(t, filepath.Join(cache, cacheFileName(testBpeURL)), 7)
|
||||
t.Setenv("TIKTOKEN_CACHE_DIR", cache)
|
||||
|
||||
ranks, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if got := ranks["hello"]; got != 7 {
|
||||
t.Errorf("rank = %d, want 7 (an explicit cache dir must win)", got)
|
||||
}
|
||||
}
|
||||
|
||||
// The whole point of this loader is that a missing table is an error rather
|
||||
// than an HTTP request. Reporting every path tried is what turns an opaque
|
||||
// "all token counts are zero" deployment into a one-look diagnosis.
|
||||
func TestLocalBpeLoader_MissingTableReportsCandidatesInsteadOfDownloading(t *testing.T) {
|
||||
isolate(t)
|
||||
|
||||
_, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err == nil {
|
||||
t.Fatal("expected an error when no local table exists, got nil")
|
||||
}
|
||||
for _, want := range []string{cacheFileName(testBpeURL), "cl100k_base.tiktoken", "TIKTOKEN_CACHE_DIR"} {
|
||||
if !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("error does not mention %q, so it cannot be acted on:\n%v", want, err)
|
||||
}
|
||||
}
|
||||
if strings.Contains(err.Error(), "http") && !strings.Contains(err.Error(), testBpeURL) {
|
||||
t.Errorf("error hints at a network attempt: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLocalBpeLoader_RejectsMalformedTable(t *testing.T) {
|
||||
dir := isolate(t)
|
||||
path := filepath.Join(dir, cacheFileName(testBpeURL))
|
||||
if err := os.WriteFile(path, []byte("not base64 at all\n"), 0o644); err != nil {
|
||||
t.Fatalf("write: %v", err)
|
||||
}
|
||||
|
||||
_, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err == nil {
|
||||
t.Fatal("expected an error for a malformed table, got nil")
|
||||
}
|
||||
}
|
||||
|
||||
// A candidate that exists but cannot be read as a file must surface as a read
|
||||
// error rather than being skipped as "missing". LoadTiktokenBpe only continues
|
||||
// past os.IsNotExist; a directory at the candidate path fails ReadFile with a
|
||||
// distinct error, which this test pins to the read-error path.
|
||||
func TestLocalBpeLoader_ReadErrorIsReported(t *testing.T) {
|
||||
dir := isolate(t)
|
||||
// A directory at the sha1-named candidate path exists but is not a
|
||||
// regular file, so os.ReadFile fails with a non-IsNotExist error.
|
||||
candidate := filepath.Join(dir, cacheFileName(testBpeURL))
|
||||
if err := os.Mkdir(candidate, 0o755); err != nil {
|
||||
t.Fatalf("mkdir: %v", err)
|
||||
}
|
||||
|
||||
_, err := localBpeLoader{}.LoadTiktokenBpe(testBpeURL)
|
||||
if err == nil {
|
||||
t.Fatal("expected a read error for a non-file candidate, got nil")
|
||||
}
|
||||
if strings.Contains(err.Error(), "no local BPE table") {
|
||||
t.Errorf("read error was masked as not-found: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -12,7 +12,7 @@
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
//go:build manual
|
||||
|
||||
package tokenizer
|
||||
|
||||
|
||||
@@ -12,7 +12,7 @@
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
//go:build manual
|
||||
|
||||
package tokenizer
|
||||
|
||||
|
||||
Reference in New Issue
Block a user