Add new embeddings

go.mod

@@ -3,8 +3,8 @@ module github.com/hupe1980/go-tiktoken
 go 1.20
 
 require (
-	github.com/dlclark/regexp2 v1.10.0
-	github.com/stretchr/testify v1.8.4
+	github.com/dlclark/regexp2 v1.11.0
+	github.com/stretchr/testify v1.9.0
 )
 
 require (

go.sum

@@ -1,11 +1,11 @@
 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.10.0 h1:+/GIL799phkJqYW+3YbOd8LCcbHzT0Pbo8zl70MHsq0=
-github.com/dlclark/regexp2 v1.10.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
+github.com/dlclark/regexp2 v1.11.0 h1:G/nrcoOa7ZXlpoa/91N3X7mM3r8eIlMBBJZvsz/mxKI=
+github.com/dlclark/regexp2 v1.11.0/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

tiktoken.go

@@ -53,6 +53,8 @@ var ModelToEncoding = map[string]string{
 	"code-davinci-edit-001": P50kEdit,
 	// embeddings
 	"text-embedding-ada-002": CL100kBase,
+	"text-embedding-3-small": CL100kBase,
+	"text-embedding-3-large": CL100kBase,
 	// old embeddings
 	"text-similarity-davinci-001":  R50kBase,
 	"text-similarity-curie-001":    R50kBase,

tiktoken_test.go

@@ -55,10 +55,13 @@ func TestNewEncodingForModel(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			encoding, err := NewEncodingForModel(tc.model)
+
 			name := ""
+
 			if encoding != nil {
 				name = encoding.Name()
 			}
+
 			assert.Equal(t, tc.expectedResult, name, "Unexpected encoding result")
 			assert.Equal(t, tc.expectedError, err, "Unexpected error")
 		})