add bpe interface

Author: pkoukk

encoding.go

@@ -116,7 +116,7 @@ func initEncoding(encodingName string) (*Encoding, error) {
 }
 
 func cl100k_base() (*Encoding, error) {
-	ranks, err := loadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken")
+	ranks, err := bpeLoader.LoadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken")
 	if err != nil {
 		return nil, err
 	}
@@ -136,7 +136,7 @@ func cl100k_base() (*Encoding, error) {
 }
 
 func p50k_edit() (*Encoding, error) {
-	ranks, err := loadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/p50k_base.tiktoken")
+	ranks, err := bpeLoader.LoadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/p50k_base.tiktoken")
 	if err != nil {
 		return nil, err
 	}
@@ -150,7 +150,7 @@ func p50k_edit() (*Encoding, error) {
 }
 
 func p50k_base() (*Encoding, error) {
-	ranks, err := loadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/p50k_base.tiktoken")
+	ranks, err := bpeLoader.LoadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/p50k_base.tiktoken")
 	if err != nil {
 		return nil, err
 	}
@@ -173,7 +173,7 @@ func p50k_base() (*Encoding, error) {
 }
 
 func r50k_base() (*Encoding, error) {
-	ranks, err := loadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/r50k_base.tiktoken")
+	ranks, err := bpeLoader.LoadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/r50k_base.tiktoken")
 	if err != nil {
 		return nil, err
 	}

go.mod

@@ -3,7 +3,7 @@ module github.com/pkoukk/tiktoken-go
 go 1.19
 
 require (
-	github.com/dlclark/regexp2 v1.8.1
+	github.com/dlclark/regexp2 v1.10.0
 	github.com/google/uuid v1.3.0
 	github.com/stretchr/testify v1.8.2
 )

go.sum

@@ -1,8 +1,8 @@
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
-github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/dlclark/regexp2 v1.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
+github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
 github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=

load.go

@@ -14,6 +14,10 @@ import (
 	"github.com/google/uuid"
 )
 
+type BpeLoader interface {
+	LoadTiktokenBpe(tiktokenBpeFile string) (map[string]int, error)
+}
+
 func readFile(blobpath string) ([]byte, error) {
 	if !strings.HasPrefix(blobpath, "http://") && !strings.HasPrefix(blobpath, "https://") {
 		file, err := os.Open(blobpath)
@@ -91,3 +95,13 @@ func loadTiktokenBpe(tiktokenBpeFile string) (map[string]int, error) {
 	}
 	return bpeRanks, nil
 }
+
+type defaultBpeLoader struct{}
+
+func (l *defaultBpeLoader) LoadTiktokenBpe(tiktokenBpeFile string) (map[string]int, error) {
+	return loadTiktokenBpe(tiktokenBpeFile)
+}
+
+func NewDefaultBpeLoader() BpeLoader {
+	return &defaultBpeLoader{}
+}

tiktoken.go

@@ -8,6 +8,12 @@ import (
 	"github.com/dlclark/regexp2"
 )
 
+var bpeLoader BpeLoader = NewDefaultBpeLoader()
+
+func SetBpeLoader(loader BpeLoader) {
+	bpeLoader = loader
+}
+
 func GetEncoding(encodingName string) (*Tiktoken, error) {
 	enc, err := getEncoding(encodingName)
 	if err != nil {