Skip to content

Commit

Permalink
tiktoken: add custom constructor
Browse files Browse the repository at this point in the history
  • Loading branch information
bobheadxi committed Sep 5, 2023
1 parent 4f571fc commit caab340
Showing 1 changed file with 10 additions and 5 deletions.
15 changes: 10 additions & 5 deletions tiktoken.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,7 @@ func GetEncoding(encodingName string) (*Tiktoken, error) {
for k := range enc.SpecialTokens {
specialTokensSet[k] = true
}
return &Tiktoken{
bpe: pbe,
pbeEncoding: enc,
specialTokensSet: specialTokensSet,
}, nil
return NewTiktoken(pbe, enc, specialTokensSet), nil
}

func EncodingForModel(modelName string) (*Tiktoken, error) {
Expand Down Expand Up @@ -121,3 +117,12 @@ func difference(setA, setB map[string]any) map[string]any {
}
return result
}

// NewTiktoken can be used to create a *Tiktoken with custom parameters.
func NewTiktoken(bpe *CoreBPE, encoding *Encoding, specialTokensSet map[string]any) *Tiktoken {
return &Tiktoken{
bpe: bpe,
pbeEncoding: encoding,
specialTokensSet: specialTokensSet,
}
}

0 comments on commit caab340

Please sign in to comment.