summaryrefslogtreecommitdiff
path: root/tokenizer/bpe/tokenizer.go
diff options
context:
space:
mode:
Diffstat (limited to 'tokenizer/bpe/tokenizer.go')
-rw-r--r--tokenizer/bpe/tokenizer.go24
1 files changed, 14 insertions, 10 deletions
diff --git a/tokenizer/bpe/tokenizer.go b/tokenizer/bpe/tokenizer.go
index 9bc351c..7e2459a 100644
--- a/tokenizer/bpe/tokenizer.go
+++ b/tokenizer/bpe/tokenizer.go
@@ -3,12 +3,14 @@ package bpe
import "github.com/jonasknobloch/mbpe"
type Tokenizer struct {
- mbpe *mbpe.Tokenizer
+ mbpe *mbpe.Tokenizer
+ config Config
}
-func NewTokenizer(mbpe *mbpe.Tokenizer) *Tokenizer {
+func NewTokenizer(mbpe *mbpe.Tokenizer, cfg Config) *Tokenizer {
return &Tokenizer{
- mbpe: mbpe,
+ mbpe: mbpe,
+ config: cfg,
}
}
@@ -29,13 +31,15 @@ func (t *Tokenizer) Decode(ids []int) string {
}
func (t *Tokenizer) Tokenize(s string) (ids []int) {
- defer func() {
- if r := recover(); r != nil {
- // defer UnknownRunes(t, s)
-
- ids = []int{}
- }
- }()
+ if t.config.Recover {
+ defer func() {
+ if r := recover(); r != nil {
+ // defer UnknownRunes(t, s)
+
+ ids = []int{}
+ }
+ }()
+ }
return t.Encode(s)
}