tiktoken-go

If you want to avoid downloading dictionaries at runtime or using the default cache, you can provide a custom BpeLoader.

To use a custom loader, call tiktoken.SetBpeLoader before calling tiktoken.GetEncoding or tiktoken.EncodingForModel. BpeLoader is an interface that you can implement yourself.

For an offline BPE loader (which loads dictionaries from embedded files), you can use the separate tiktoken_loader project: github.com/pkoukk/tiktoken-go-loader.

// Example of setting an offline loader
// tiktoken.SetBpeLoader(tiktoken_loader.NewOfflineLoader())

By default, tiktoken-go downloads dictionaries at runtime. To avoid runtime downloads or to use a local cache, you can implement the BpeLoader interface and set it before calling tiktoken.GetEncoding or tiktoken.EncodingForModel using tiktoken.SetBpeLoader.

For an offline BPE loader that loads dictionaries from embedded files, you can use the tiktoken_loader package: github.com/pkoukk/tiktoken-go-loader.

// Example of setting an offline loader
// tiktoken.SetBpeLoader(tiktoken_loader.NewOfflineLoader())

To use tiktoken-go in your Go project, install it using go get:

go get github.com/pkoukk/tiktoken-go

The library provides internal mapping between OpenAI model names and their required BPE encodings. This allows you to resolve which encoding to use based on the model you are targeting.

Model to Encoding Mappings:

• gpt-4.5, gpt-4.1, gpt-4o $\rightarrow$ o200k_base
• gpt-4, gpt-3.5-turbo, text-embedding-3-large $\rightarrow$ cl100k_base
• text-davinci-003, code-davinci-002 $\rightarrow$ p50k_base
• text-davinci-001, davinci, curie $\rightarrow$ r50k_base
• text-davinci-edit-001 $\rightarrow$ p50k_edit
• gpt2 $\rightarrow$ gpt2

Prefix-based Mappings: If you have a specific versioned model name (e.g., gpt-4o-2024-05-13), the library uses MODEL_PREFIX_TO_ENCODING to match the prefix (e.g., gpt-4o-) to the correct base encoding.

tiktoken-go uses a cache mechanism for token dictionaries. You can control where these are stored using the TIKTOKEN_CACHE_DIR environment variable.

• If set: The library will use the specified directory to cache the dictionary.
• If not set: The library will download the dictionary from the internet every time you initialize an encoding for the first time.

Like the original OpenAI tiktoken library, tiktoken-go supports a caching mechanism for token dictionaries.

You can set the cache directory using the TIKTOKEN_CACHE_DIR environment variable. If this variable is set, the library will use that directory to cache dictionaries. If not set, tiktoken-go will download the dictionaries every time an encoding is initialized for the first time.

Use tiktoken.GetEncoding(encoding) to retrieve an encoding by its specific name (e.g., cl100k_base). Once you have the encoding object, use .Encode(text, nil, nil) to convert text into tokens.

package main

import (
    "fmt"
    "github.com/pkoukk/tiktoken-go"
)

func main()  {
	text := "Hello, world!"
	encoding := "cl100k_base"

	ke, err := tiktoken.GetEncoding(encoding)
	if err != nil {
		err = fmt.Errorf("getEncoding: %v", err)
		return
	}

	// encode
	token := tke.Encode(text, nil, nil)

	// tokens
	fmt.Println((token))
	// num_tokens
	fmt.Println(len(token))
}

Use tiktoken.EncodingForModel(model) to retrieve an encoding instance associated with a specific OpenAI model name (e.g., gpt-3.5-turbo). This is often more convenient than manually looking up the encoding type.

package main

import (
	"fmt"
	"github.com/pkoukk/tiktoken-go"
)

func main() {
	text := "Hello, world!"
	encoding := "gpt-3.5-turbo"

	tkm, err := tiktoken.EncodingForModel(encoding)
	if err != nil {
		fmt.Printf("getEncoding: %v\n", err)
		return
	}

	// encode
	token := tkm.Encode(text, nil, nil)

	// tokens
	fmt.Println(token)
	// num_tokens
	fmt.Println(len(token))
}

The defaultBpeLoader uses a cache to store downloaded BPE files. You can control where these files are stored by setting one of the following environment variables in your system:

1. TIKTOKEN_CACHE_DIR: The primary directory used for caching.
2. DATA_GYM_CACHE_DIR: A fallback directory if TIKTOKEN_CACHE_DIR is not set.

If neither environment variable is provided, the loader defaults to using a directory named data-gym-cache within your system's temporary directory (os.TempDir()).

The following encoding names are supported and map to specific OpenAI models:

| Encoding name           | OpenAI models                                                                                          |
|-------------------------|--------------------------------------------------------------------------------------------------------|
| `o200k_base`            | `gpt-4o`, `gpt-4.1`, `gpt-4.5`                                                                         |
| `cl100k_base`           | `gpt-4`, `gpt-3.5-turbo`, `text-embedding-ada-002`, `text-embedding-3-small`, `text-embedding-3-large` |
| `p50k_base`             | Codex models, `text-davinci-002`, `text-davinci-003`                                                    |
| `r50k_base` (or `gpt2`) | GPT-3 models like `davinci`                                                                             |