Handle unknown runes

3 files changed, 119 insertions, 1 deletions

diff --git a/tokenizer/bpe/cmd/tokenize/main.go b/tokenizer/bpe/cmd/tokenize/main.go
index 1b0b868..1dbfaae 100644
--- a/tokenizer/bpe/cmd/tokenize/main.go
+++ b/tokenizer/bpe/cmd/tokenize/main.go
@@ -109,5 +109,7 @@ func tokenizer() llm.Tokenizer {
 		tok = t
 	}
 
+	_ = bpe.ByteCoverage(tok)
+
 	return tok
 }
diff --git a/tokenizer/bpe/tokenizer.go b/tokenizer/bpe/tokenizer.go
index 1b65de2..39dd7f4 100644
--- a/tokenizer/bpe/tokenizer.go
+++ b/tokenizer/bpe/tokenizer.go
@@ -20,6 +20,14 @@ func (t *Tokenizer) Decode(ids []int) string {
 	panic("unimplemented") // TODO implement
 }
 
-func (t *Tokenizer) Tokenize(s string) []int {
+func (t *Tokenizer) Tokenize(s string) (ids []int) {
+	defer func() {
+		if r := recover(); r != nil {
+			// defer UnknownRunes(t, s)
+
+			ids = []int{}
+		}
+	}()
+
 	return t.Encode(s)
 }
diff --git a/tokenizer/bpe/validate.go b/tokenizer/bpe/validate.go
new file mode 100644
index 0000000..f9bd265
--- /dev/null
+++ b/tokenizer/bpe/validate.go
@@ -0,0 +1,108 @@
+package bpe
+
+import (
+	"fmt"
+	"slices"
+
+	"github.com/jonasknobloch/mbpe"
+)
+
+func InitialAlphabet() []rune {
+	alphabet := make([]rune, 256)
+
+	bc := mbpe.BytesChar
+
+	for i := 0; i < 256; i++ {
+		b := uint8(i)
+
+		runes := []rune(bc[b])
+
+		if len(runes) != 1 {
+			panic("unexpected replacement")
+		}
+
+		alphabet[i] = runes[0]
+	}
+
+	slices.Sort(alphabet)
+
+	return alphabet
+}
+
+func UnknownRunes(t *Tokenizer, s string) []rune {
+	atoi := make(map[string]int)
+
+	vocab := t.mbpe.Model().(*mbpe.MBPE).Vocab()
+
+	for i, token := range vocab {
+		if _, ok := atoi[token]; ok {
+			continue
+		}
+
+		atoi[token] = i
+	}
+
+	unknown := make(map[rune]struct{})
+
+	chunks := t.mbpe.PreTokenizer().PreTokenize(s)
+
+	for _, chunk := range chunks {
+		for _, r := range chunk {
+			i, ok := atoi[string(r)]
+
+			if !ok {
+				fmt.Printf("%d %s %v not in vocabulary\n", i, string(r), []byte(string(r)))
+
+				if _, ok := unknown[r]; ok {
+					continue
+				}
+
+				unknown[r] = struct{}{}
+			}
+		}
+	}
+
+	result := make([]rune, 0, len(unknown))
+
+	for r := range unknown {
+		result = append(result, r)
+	}
+
+	slices.Sort(result)
+
+	return result
+}
+
+func ByteCoverage(t *Tokenizer) bool {
+	atoi := make(map[string]int)
+
+	vocab := t.mbpe.Model().(*mbpe.MBPE).Vocab()
+
+	for i, token := range vocab {
+		if _, ok := atoi[token]; ok {
+			continue
+		}
+
+		atoi[token] = i
+	}
+
+	bc := mbpe.BytesChar
+
+	covered := true
+
+	for i := 0; i < 256; i++ {
+		c, ok := bc[byte(i)]
+
+		if !ok {
+			panic("not in replacement table")
+		}
+
+		if _, ok := atoi[c]; !ok {
+			fmt.Printf("%d %s %v not in vocabulary\n", i, c, []byte(c))
+
+			covered = false
+		}
+	}
+
+	return covered
+}