diff --git a/go.mod b/go.mod index 8982bf9f70b1..abd72725168c 100644 --- a/go.mod +++ b/go.mod @@ -458,7 +458,7 @@ require ( github.com/peterbourgon/diskv v2.0.1+incompatible // indirect github.com/pierrec/lz4/v4 v4.1.8 // indirect github.com/pkg/errors v0.9.1 - github.com/pkoukk/tiktoken-go v0.1.7 // indirect + github.com/pkoukk/tiktoken-go v0.1.7 github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect github.com/polydawn/refmt v0.89.1-0.20231129105047-37766d95467a // indirect github.com/power-devops/perfstat v0.0.0-20240221224432-82ca36839d55 // indirect diff --git a/pkg/tokens/tokens.go b/pkg/tokens/tokens.go new file mode 100644 index 000000000000..2a41b64fa844 --- /dev/null +++ b/pkg/tokens/tokens.go @@ -0,0 +1,103 @@ +// SPDX-License-Identifier: MIT + +package tokens + +import ( + "fmt" + "strings" + + "github.com/mudler/LocalAI/core/schema" + tiktoken "github.com/pkoukk/tiktoken-go" +) + +const defaultEncoding = tiktoken.MODEL_CL100K_BASE + +// EncodingFor resolves the encoding known by tiktoken and uses cl100k_base +// when the supplied model has no known mapping. +func EncodingFor(model string) string { + if encoding, ok := tiktoken.MODEL_TO_ENCODING[model]; ok { + return encoding + } + for prefix, encoding := range tiktoken.MODEL_PREFIX_TO_ENCODING { + if strings.HasPrefix(model, prefix) { + return encoding + } + } + return defaultEncoding +} + +// CountText counts text tokens without interpreting literal special-token +// strings as tokenizer control tokens. +func CountText(text, model string) (int, error) { + encoding, err := tiktoken.GetEncoding(EncodingFor(model)) + if err != nil { + return 0, fmt.Errorf("get token encoding: %w", err) + } + return len(encoding.EncodeOrdinary(text)), nil +} + +// Count counts only textual message content. +func Count(messages []schema.Message, model string) (int, error) { + total := 0 + for index, message := range messages { + text, err := messageText(message.Content) + if err != nil { + return 0, fmt.Errorf("message %d: %w", index, err) + } + count, err := CountText(text, model) + if err != nil { + return 0, err + } + total += count + } + return total, nil +} + +func messageText(content any) (string, error) { + switch value := content.(type) { + case nil: + return "", nil + case string: + return value, nil + case []any: + var text strings.Builder + for index, rawPart := range value { + part, ok := rawPart.(map[string]any) + if !ok { + return "", fmt.Errorf("unsupported multimodal content part %d of type %T", index, rawPart) + } + + if rawText, exists := part["text"]; exists { + partText, ok := rawText.(string) + if !ok { + return "", fmt.Errorf("unsupported multimodal text part %d: text has type %T", index, rawText) + } + text.WriteString(partText) + continue + } + + partType, ok := part["type"].(string) + if ok && recognizedMediaType(partType) { + continue + } + return "", fmt.Errorf("unsupported multimodal content part %d", index) + } + return text.String(), nil + case map[string]any: + if len(value) == 0 { + return "", nil + } + return "", fmt.Errorf("unsupported message content of type %T", content) + default: + return "", fmt.Errorf("unsupported message content of type %T", content) + } +} + +func recognizedMediaType(partType string) bool { + switch partType { + case "image", "image_url", "audio", "audio_url", "input_audio", "video", "video_url": + return true + default: + return false + } +} diff --git a/pkg/tokens/tokens_suite_test.go b/pkg/tokens/tokens_suite_test.go new file mode 100644 index 000000000000..22ef9d0be0c0 --- /dev/null +++ b/pkg/tokens/tokens_suite_test.go @@ -0,0 +1,34 @@ +// SPDX-License-Identifier: MIT + +package tokens_test + +import ( + "testing" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + tiktoken "github.com/pkoukk/tiktoken-go" +) + +type testBPELoader struct{} + +func (testBPELoader) LoadTiktokenBpe(string) (map[string]int, error) { + ranks := make(map[string]int, 256) + for value := 0; value < 256; value++ { + ranks[string([]byte{byte(value)})] = value + } + for index, token := range []string{ + "he", "hel", "hell", "hello", + "wo", "wor", "worl", "world", + " w", " wo", " wor", " worl", " world", + } { + ranks[token] = 256 + index + } + return ranks, nil +} + +func TestTokens(t *testing.T) { + RegisterFailHandler(Fail) + tiktoken.SetBpeLoader(testBPELoader{}) + RunSpecs(t, "Tokens Suite") +} diff --git a/pkg/tokens/tokens_test.go b/pkg/tokens/tokens_test.go new file mode 100644 index 000000000000..fec9ac997848 --- /dev/null +++ b/pkg/tokens/tokens_test.go @@ -0,0 +1,79 @@ +// SPDX-License-Identifier: MIT + +package tokens_test + +import ( + "github.com/mudler/LocalAI/core/schema" + "github.com/mudler/LocalAI/pkg/tokens" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("Token counting", func() { + Describe("EncodingFor", func() { + It("uses tiktoken's exact model mapping", func() { + Expect(tokens.EncodingFor("text-davinci-003")).To(Equal("p50k_base")) + }) + + It("uses tiktoken's model prefix mapping", func() { + Expect(tokens.EncodingFor("gpt-4o-2024-08-06")).To(Equal("o200k_base")) + }) + + It("falls back conservatively for unknown models", func() { + Expect(tokens.EncodingFor("local-model")).To(Equal("cl100k_base")) + }) + }) + + Describe("CountText", func() { + It("returns the exact encoded token count", func() { + count, err := tokens.CountText("hello world", "gpt-4") + Expect(err).NotTo(HaveOccurred()) + Expect(count).To(Equal(2)) + }) + + It("treats literal special-token strings as ordinary text", func() { + var count int + Expect(func() { + var err error + count, err = tokens.CountText("<|endoftext|>", "gpt-4") + Expect(err).NotTo(HaveOccurred()) + }).NotTo(Panic()) + Expect(count).To(BeNumerically(">", 0)) + }) + }) + + Describe("Count", func() { + It("sums string message content", func() { + count, err := tokens.Count([]schema.Message{ + {Content: "hello"}, + {Content: "world"}, + }, "gpt-4") + Expect(err).NotTo(HaveOccurred()) + Expect(count).To(Equal(2)) + }) + + It("extracts text from multimodal content and ignores recognized media", func() { + count, err := tokens.Count([]schema.Message{{ + Content: []any{ + map[string]any{"type": "text", "text": "hello world"}, + map[string]any{"type": "image_url", "image_url": map[string]any{"url": "image.png"}}, + }, + }}, "gpt-4") + Expect(err).NotTo(HaveOccurred()) + Expect(count).To(Equal(2)) + }) + + DescribeTable("rejects unsupported non-empty content", + func(content any) { + _, err := tokens.Count([]schema.Message{{Content: content}}, "gpt-4") + Expect(err).To(MatchError(ContainSubstring("unsupported"))) + }, + Entry("top-level object", map[string]any{"text": "hello"}), + Entry("arbitrary multimodal map", []any{map[string]any{"foo": "bar"}}), + Entry("malformed text part", []any{map[string]any{"type": "text", "text": 42}}), + Entry("unknown typed part", []any{map[string]any{"type": "custom", "custom": "value"}}), + Entry("non-object multimodal part", []any{"hello"}), + ) + }) +})