Expose finite state automaton

1 files changed, 44 insertions, 0 deletions

diff --git a/tokenizer/bpe/bytelevel.go b/tokenizer/bpe/bytelevel.go
new file mode 100644
index 0000000..4059d0d
--- /dev/null
+++ b/tokenizer/bpe/bytelevel.go
@@ -0,0 +1,44 @@
+package bpe
+
+import (
+	"github.com/jonasknobloch/mbpe"
+	"go.jknobloc.com/x/tokenizer/bpe/split"
+)
+
+type ByteLevel struct {
+	*mbpe.ByteLevel
+	addPrefixSpace bool
+	fsa            *split.FSA
+}
+
+func NewByteLevel(addPrefixSpace bool) *ByteLevel {
+	return &ByteLevel{
+		ByteLevel:      mbpe.NewByteLevel(addPrefixSpace),
+		addPrefixSpace: addPrefixSpace,
+		fsa:            split.NewFSA(),
+	}
+}
+
+func (p *ByteLevel) PreTokenize(phrase string) []string {
+	if phrase == "" {
+		return []string{}
+	}
+
+	if p.addPrefixSpace && phrase[0] != ' ' {
+		phrase = " " + phrase
+	}
+
+	compounds := p.fsa.FindAll(phrase)
+
+	for i, compound := range compounds {
+		r := ""
+
+		for _, b := range []byte(compound) {
+			r += mbpe.BytesChar[b]
+		}
+
+		compounds[i] = r
+	}
+
+	return compounds
+}