File size: 741 Bytes
8d3471e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 | package util
const (
defaultTokenizerModel = "gpt-4o"
claudeTokenizerModel = "claude"
)
func CountPromptTokens(text, model string) int {
base := maxTokenCount(
EstimateTokens(text),
countWithTokenizer(text, model),
)
if base <= 0 {
return 0
}
return base + conservativePromptPadding(base)
}
func CountOutputTokens(text, model string) int {
base := maxTokenCount(
EstimateTokens(text),
countWithTokenizer(text, model),
)
if base <= 0 {
return 0
}
return base
}
func conservativePromptPadding(base int) int {
padding := base / 50
if padding < 4 {
padding = 4
}
return padding
}
func maxTokenCount(values ...int) int {
best := 0
for _, v := range values {
if v > best {
best = v
}
}
return best
}
|