103 lines
2.6 KiB
Go
103 lines
2.6 KiB
Go
package tokenkit
|
|
|
|
import (
|
|
"strings"
|
|
"sync"
|
|
"unicode"
|
|
"unicode/utf8"
|
|
|
|
"github.com/tiktoken-go/tokenizer"
|
|
"github.com/tiktoken-go/tokenizer/codec"
|
|
)
|
|
|
|
var (
|
|
o200kBase = sync.OnceValue(func() tokenizer.Codec { return codec.NewO200kBase() })
|
|
cl100kBase = sync.OnceValue(func() tokenizer.Codec { return codec.NewCl100kBase() })
|
|
)
|
|
|
|
// maxRunBytes bounds a run of letters, whitespace, or symbols passed to the
|
|
// encoder in one piece. BPE cost grows with the square of a pre-tokenized
|
|
// piece: 100,000 Chinese characters without punctuation take about 20 s in one
|
|
// piece. Cutting longer runs keeps counting linear; ordinary text has no such
|
|
// runs, and a cut changes the count only inside the run (under 0.5% for CJK).
|
|
const maxRunBytes = 1024
|
|
|
|
// Count returns the token count of text for model. GPT and o-series models
|
|
// are counted with their tiktoken encoding; other models use Estimate.
|
|
func Count(model, text string) int {
|
|
if text == "" {
|
|
return 0
|
|
}
|
|
model = strings.ToLower(model)
|
|
if ModelFamily(model) != FamilyOpenAI {
|
|
return Estimate(model, text)
|
|
}
|
|
encoder := o200kBase()
|
|
if usesCl100k(model) {
|
|
encoder = cl100kBase()
|
|
}
|
|
|
|
total, start := 0, 0
|
|
runClass, runStart := runNone, 0
|
|
for i, r := range text {
|
|
class := runClassOf(r)
|
|
if class != runClass {
|
|
runClass, runStart = class, i
|
|
continue
|
|
}
|
|
if class != runNumber && i-runStart >= maxRunBytes {
|
|
n, _ := encoder.Count(text[start:i])
|
|
total += n
|
|
start, runStart = i, i
|
|
}
|
|
}
|
|
n, _ := encoder.Count(text[start:])
|
|
return total + n
|
|
}
|
|
|
|
type runClass int
|
|
|
|
const (
|
|
runNone runClass = iota
|
|
runLetter
|
|
runNumber
|
|
runSpace
|
|
runSymbol
|
|
)
|
|
|
|
// runClassOf groups runes the way tiktoken's pre-tokenizer joins them into
|
|
// pieces. Digits are split into groups of three by the encoder itself.
|
|
func runClassOf(r rune) runClass {
|
|
if r < utf8.RuneSelf {
|
|
switch {
|
|
case 'a' <= r && r <= 'z' || 'A' <= r && r <= 'Z':
|
|
return runLetter
|
|
case '0' <= r && r <= '9':
|
|
return runNumber
|
|
case r == ' ' || '\t' <= r && r <= '\r':
|
|
return runSpace
|
|
default:
|
|
return runSymbol
|
|
}
|
|
}
|
|
switch {
|
|
case unicode.IsLetter(r) || unicode.IsMark(r):
|
|
return runLetter
|
|
case unicode.IsNumber(r):
|
|
return runNumber
|
|
case unicode.IsSpace(r):
|
|
return runSpace
|
|
default:
|
|
return runSymbol
|
|
}
|
|
}
|
|
|
|
// usesCl100k reports whether an OpenAI model predates o200k_base: GPT-4 (but
|
|
// not GPT-4o, GPT-4.1, or GPT-4.5) and GPT-3.5. Every later model, including
|
|
// GPT-5, GPT-6, and the o-series, uses o200k_base.
|
|
func usesCl100k(model string) bool {
|
|
model = strings.TrimPrefix(model, "ft:")
|
|
return model == "gpt-4" || strings.HasPrefix(model, "gpt-4-") ||
|
|
strings.HasPrefix(model, "gpt-3.5") || strings.HasPrefix(model, "gpt-35")
|
|
}
|