52 lines
994 B
Go
52 lines
994 B
Go
package tokens
|
|
|
|
import "github.com/pkoukk/tiktoken-go"
|
|
|
|
var tke *tiktoken.Tiktoken
|
|
|
|
func getEncoding() *tiktoken.Tiktoken {
|
|
if tke == nil {
|
|
t, err := tiktoken.GetEncoding("o200k_base")
|
|
if err != nil {
|
|
return nil
|
|
}
|
|
tke = t
|
|
}
|
|
return tke
|
|
}
|
|
|
|
// Tokenize 将文本编码为 o200k_base token id 列表（去重）。
|
|
// 编码器初始化失败时返回 nil。
|
|
func Tokenize(text string) []int64 {
|
|
if text == "" {
|
|
return nil
|
|
}
|
|
enc := getEncoding()
|
|
if enc == nil {
|
|
return nil
|
|
}
|
|
seen := make(map[int64]bool)
|
|
var out []int64
|
|
for _, id := range enc.Encode(text, nil, nil) {
|
|
tid := int64(id)
|
|
if seen[tid] {
|
|
continue
|
|
}
|
|
seen[tid] = true
|
|
out = append(out, tid)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Count 统计文本的 token 数量；编码器初始化失败时回退为字符数/2 估算。
|
|
func Count(text string) int64 {
|
|
if text == "" {
|
|
return 0
|
|
}
|
|
enc := getEncoding()
|
|
if enc == nil {
|
|
return int64(len([]rune(text)) / 2)
|
|
}
|
|
return int64(len(enc.Encode(text, nil, nil)))
|
|
}
|