estimate tokens

tiktoken.go

@@ -0,0 +1,150 @@
+package main
+
+import (
+	"bufio"
+	"encoding/base64"
+	"errors"
+	"fmt"
+	"net/http"
+	"strconv"
+	"strings"
+)
+
+const TikTokenSource = "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken"
+
+type TreeNode struct {
+	TokenID  int
+	Children map[byte]*TreeNode
+}
+
+type Tokenizer struct {
+	Root *TreeNode
+}
+
+func NewTreeNode() *TreeNode {
+	return &TreeNode{
+		TokenID:  -1,
+		Children: make(map[byte]*TreeNode),
+	}
+}
+
+func (n *TreeNode) Insert(token []byte, id int) {
+	curr := n
+
+	for _, b := range token {
+		if _, ok := curr.Children[b]; !ok {
+			curr.Children[b] = NewTreeNode()
+		}
+
+		curr = curr.Children[b]
+	}
+
+	curr.TokenID = id
+}
+
+func LoadTokenizer(url string) (*Tokenizer, error) {
+	log.Println("Loading tokenizer...")
+
+	vocabulary, err := LoadVocabulary(url)
+	if err != nil {
+		return nil, err
+	}
+
+	root := NewTreeNode()
+
+	for tokenStr, id := range vocabulary {
+		root.Insert([]byte(tokenStr), id)
+	}
+
+	return &Tokenizer{
+		Root: root,
+	}, nil
+}
+
+func (t *Tokenizer) Encode(text string) []int {
+	var (
+		index  int
+		tokens []int
+	)
+
+	input := []byte(text)
+
+	for index < len(input) {
+		bestMatchLength := 0
+		bestMatchID := -1
+
+		currNode := t.Root
+
+		for i := index; i < len(input); i++ {
+			b := input[i]
+
+			childNode, exists := currNode.Children[b]
+			if !exists {
+				break
+			}
+
+			currNode = childNode
+
+			if currNode.TokenID != -1 {
+				bestMatchID = currNode.TokenID
+				bestMatchLength = (i - index) + 1
+			}
+		}
+
+		// should not be possible
+		if bestMatchLength == 0 {
+			bestMatchLength = 1
+		}
+
+		if bestMatchID != -1 {
+			tokens = append(tokens, bestMatchID)
+		}
+
+		index += bestMatchLength
+	}
+
+	return tokens
+}
+
+func LoadVocabulary(url string) (map[string]int, error) {
+	resp, err := http.Get(url)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, errors.New(resp.Status)
+	}
+
+	vocab := make(map[string]int)
+
+	scanner := bufio.NewScanner(resp.Body)
+
+	for scanner.Scan() {
+		parts := strings.SplitN(scanner.Text(), " ", 2)
+
+		if len(parts) != 2 {
+			continue
+		}
+
+		decoded, err := base64.StdEncoding.DecodeString(parts[0])
+		if err != nil {
+			return nil, fmt.Errorf("failed to decode token '%s': %w", parts[0], err)
+		}
+
+		id, err := strconv.Atoi(parts[1])
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse token ID '%s': %w", parts[1], err)
+		}
+
+		vocab[string(decoded)] = id
+	}
+
+	if err := scanner.Err(); err != nil {
+		return nil, err
+	}
+
+	return vocab, nil
+}