84 lines
1.6 KiB
Go
84 lines
1.6 KiB
Go
package tokens
|
||
|
||
import (
|
||
"strings"
|
||
|
||
"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
|
||
}
|
||
|
||
// Token 表示一个 o200k_base token：ID 与对应的符号文本。
|
||
type Token struct {
|
||
ID int64
|
||
Text string
|
||
}
|
||
|
||
// Tokenize 将文本编码为 token 列表（去重），Text 为 token 对应的符号。
|
||
// 文本先统一转小写，使索引与查询大小写不敏感；编码器初始化失败时返回 nil。
|
||
func Tokenize(text string) []Token {
|
||
if text == "" {
|
||
return nil
|
||
}
|
||
text = strings.ToLower(text)
|
||
enc := getEncoding()
|
||
if enc == nil {
|
||
return nil
|
||
}
|
||
seen := make(map[int64]bool)
|
||
var out []Token
|
||
for _, id := range enc.Encode(text, nil, nil) {
|
||
tid := int64(id)
|
||
if seen[tid] {
|
||
continue
|
||
}
|
||
seen[tid] = true
|
||
out = append(out, Token{ID: tid, Text: enc.Decode([]int{int(tid)})})
|
||
}
|
||
return out
|
||
}
|
||
|
||
// TokenIDs 提取 Token 列表的 id。
|
||
func TokenIDs(ts []Token) []int64 {
|
||
if len(ts) == 0 {
|
||
return nil
|
||
}
|
||
out := make([]int64, len(ts))
|
||
for i, t := range ts {
|
||
out[i] = t.ID
|
||
}
|
||
return out
|
||
}
|
||
|
||
// Text 返回 token id 对应的符号文本；未知 id 返回空串。
|
||
func Text(id int64) string {
|
||
enc := getEncoding()
|
||
if enc == nil {
|
||
return ""
|
||
}
|
||
return enc.Decode([]int{int(id)})
|
||
}
|
||
|
||
// 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)))
|
||
}
|