55 lines
1.6 KiB
Go
55 lines
1.6 KiB
Go
package chunking
|
|
|
|
import (
|
|
"unicode"
|
|
)
|
|
|
|
// TokenEstimator provides a deterministic token estimation suitable for prompt budgeting.
|
|
// The estimator is approximate but stable, isolated, and replaceable.
|
|
type TokenEstimator interface {
|
|
EstimateTokens(text string) int
|
|
}
|
|
|
|
// SimpleTokenEstimator provides a basic deterministic token estimation.
|
|
// This uses a simple heuristic based on word count and punctuation.
|
|
type SimpleTokenEstimator struct{}
|
|
|
|
// NewSimpleTokenEstimator creates a new simple token estimator.
|
|
func NewSimpleTokenEstimator() *SimpleTokenEstimator {
|
|
return &SimpleTokenEstimator{}
|
|
}
|
|
|
|
// EstimateTokens provides a rough estimate of the number of tokens in the given text.
|
|
// This implementation uses a simple heuristic: count words and punctuation as tokens.
|
|
// The estimate is deterministic and stable for the same input text.
|
|
func (e *SimpleTokenEstimator) EstimateTokens(text string) int {
|
|
words, punctuation := e.CountComponents(text)
|
|
return words + punctuation/2
|
|
}
|
|
|
|
// CountComponents returns unrounded counts for budgets composed from fragments.
|
|
func (e *SimpleTokenEstimator) CountComponents(text string) (words, punctuation int) {
|
|
inWord := false
|
|
for _, r := range text {
|
|
nonspace := !unicode.IsSpace(r)
|
|
if nonspace && !inWord {
|
|
words++
|
|
}
|
|
inWord = nonspace
|
|
if unicode.IsPunct(r) && r != '\'' && r != '-' && r != '_' {
|
|
punctuation++
|
|
}
|
|
}
|
|
return words, punctuation
|
|
}
|
|
|
|
// ConstTokenEstimator returns a constant token count for testing purposes.
|
|
type ConstTokenEstimator struct {
|
|
Tokens int
|
|
}
|
|
|
|
// EstimateTokens returns the configured constant token count.
|
|
func (e *ConstTokenEstimator) EstimateTokens(text string) int {
|
|
return e.Tokens
|
|
}
|