Use go:embed for pretrained gpt2 tokenizer, add missing json

.gitignore

@@ -8,4 +8,3 @@ example/testdata/
 *.dat
 *.log
 
-*.json

go.mod

@@ -1,13 +1,17 @@
 module github.com/sugarme/tokenizer
 
-go 1.13
+go 1.18
 
 require (
 	github.com/emirpasic/gods v1.12.0
 	github.com/rivo/uniseg v0.1.0
 	github.com/schollz/progressbar/v2 v2.15.0
-	github.com/stretchr/testify v1.4.0 // indirect
 	github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c
 	golang.org/x/text v0.3.3
+)
+
+require (
+	github.com/mitchellh/colorstring v0.0.0-20190213212951-d06e56a500db // indirect
+	github.com/stretchr/testify v1.4.0 // indirect
 	gopkg.in/yaml.v2 v2.2.7 // indirect
 )

go.sum

@@ -12,7 +12,6 @@ github.com/rivo/uniseg v0.1.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJ
 github.com/schollz/progressbar/v2 v2.15.0 h1:dVzHQ8fHRmtPjD3K10jT3Qgn/+H+92jhPrhmxIJfDz8=
 github.com/schollz/progressbar/v2 v2.15.0/go.mod h1:UdPq3prGkfQ7MOzZKlDRpYKcFqEMczbD7YmbPgpzKMI=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
-github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
 github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
 github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
@@ -21,9 +20,7 @@ github.com/sugarme/regexpset v0.0.0-20200920021344-4d4ec8eaf93c/go.mod h1:2gwkXL
 golang.org/x/text v0.3.3 h1:cokOdA+Jmi5PJGXLlLllQSgYigAEfHXJAERHVMaCc2k=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
-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.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.7 h1:VUgggvou5XRW9mHwD/yXxIYSMtY0zoKQf/v226p2nyo=
 gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

model/bpe/bpe.go

@@ -2,14 +2,17 @@ package bpe
 
 import (
 	"bufio"
+	"embed"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"path/filepath"
 	"regexp"
 	"sort"
+
 	// "strconv"
 	"log"
 	"strings"
@@ -35,6 +38,7 @@ type Config struct {
 	unkToken                *string
 	continuingSubwordPrefix *string
 	endOfWordSuffix         *string
+	fs                      *embed.FS
 }
 
 // BpeBuilder can be used to create a `BPE` model with
@@ -58,6 +62,7 @@ func NewBpeBuilder() *BpeBuilder {
 			unkToken:                nil,
 			continuingSubwordPrefix: nil,
 			endOfWordSuffix:         nil,
+			fs:                      nil,
 		},
 	}
 }
@@ -67,6 +72,11 @@ func (bb *BpeBuilder) Files(vocab string, merges string) {
 	bb.config.files = &configFiles{vocab, merges}
 }
 
+// FS sets the fs to use to load embedded vocab and merges files.
+func (bb *BpeBuilder) FS(fs embed.FS) {
+	bb.config.fs = &fs
+}
+
 // VocabAndMerges sets vocab and merges
 func (bb *BpeBuilder) VocabAndMerges(vocab model.Vocab, merges Merges) {
 	bb.config.vocab = &vocab
@@ -123,6 +133,11 @@ func (bb *BpeBuilder) Build() (*BPE, error) {
 		}
 	}
 
+	// Have the BPE read the files from |bb.config.fs| if set.
+	if bb.config.fs != nil {
+		bpe.fs = bb.config.fs
+	}
+
 	// Read files if provided
 	if bb.config.files != nil {
 		vocab, merges, err = bpe.ReadFiles(bb.config.files.vocab, bb.config.files.merges)
@@ -190,6 +205,10 @@ type BPE struct {
 	// EndOfWordSuffix is an optional suffix
 	// to caracterize and end-of-word subword
 	EndOfWordSuffix *string
+
+	// fs is an embedded file system. It allows you to import vocab/merges files using go:embed.
+	// If |fs| is not nil, it is preferred over the local filesystem.
+	fs *embed.FS
 }
 
 func (b *BPE) builder() *BpeBuilder {
@@ -230,6 +249,14 @@ func NewBpeFromFiles(vocab, merges string) (*BPE, error) {
 	return b.Build()
 }
 
+// NewBPEFromFS create BPE model from |vocab| and |merges| files in |fs|.
+func NewBPEFromFS(fs embed.FS, vocab, merges string) (*BPE, error) {
+	b := NewBpeBuilder()
+	b.FS(fs)
+	b.Files(vocab, merges)
+	return b.Build()
+}
+
 // NewBPE creates new BPE model with given vocab and merges
 func NewBPE(vocab model.Vocab, merges Merges) *BPE {
 	b, err := newBPE()
@@ -260,7 +287,15 @@ func (b *BPE) FromFiles(vocab string, merges string) *BpeBuilder {
 func (b *BPE) ReadFiles(vocabF string, mergesF string) (*model.Vocab, *Merges, error) {
 	var err error
 	// read json file
-	vocabBytes, err := ioutil.ReadFile(vocabF)
+
+	var vocabBytes []byte
+
+	if b.fs != nil {
+		vocabBytes, err = b.fs.ReadFile(vocabF)
+	} else {
+		vocabBytes, err = ioutil.ReadFile(vocabF)
+	}
+
 	if err != nil {
 		return nil, nil, err
 	}
@@ -275,9 +310,15 @@ func (b *BPE) ReadFiles(vocabF string, mergesF string) (*model.Vocab, *Merges, e
 		return nil, nil, err
 	}
 
+	var mFile io.ReadCloser
 	// Read merges file. Each line contains a Merges object(rank, )
 	// Recall: Merges is map[Pair]PairVal (rank int, newId int)
-	mFile, err := os.Open(mergesF)
+	if b.fs != nil {
+		mFile, err = b.fs.Open(mergesF)
+	} else {
+		mFile, err = os.Open(mergesF)
+	}
+
 	if err != nil {
 		return nil, nil, err
 	}

pretrained/gpt2.go

@@ -1,15 +1,14 @@
 package pretrained
 
 import (
+	"embed"
 	"log"
-	"os"
 
 	"github.com/sugarme/tokenizer"
 	"github.com/sugarme/tokenizer/decoder"
 	"github.com/sugarme/tokenizer/model/bpe"
 	"github.com/sugarme/tokenizer/pretokenizer"
 	"github.com/sugarme/tokenizer/processor"
-	"github.com/sugarme/tokenizer/util"
 )
 
 // GPT2 loads GPT2 (small) tokenizer from vocab and merges files.
@@ -29,19 +28,18 @@ import (
 // Source:
 // "https://cdn.huggingface.co/gpt2-merges.txt"
 // "https://cdn.huggingface.co/gpt2-vocab.json"
-func GPT2(addPrefixSpace bool, trimOffsets bool) *tokenizer.Tokenizer {
 
-	currDir, err := os.Getwd()
-	if err != nil {
-		log.Fatal(err)
-	}
-	util.CdToThis()
-	defer util.CdBack(currDir)
+//go:embed model/gpt2-merges.txt
+//go:embed model/gpt2-vocab.json
+var fs embed.FS
 
-	vocabFile := "model/gpt2-vocab.json"
-	mergeFile := "model/gpt2-merges.txt"
+const (
+	vocabFilename = "model/gpt2-vocab.json"
+	mergeFilename = "model/gpt2-merges.txt"
+)
 
-	model, err := bpe.NewBpeFromFiles(vocabFile, mergeFile)
+func GPT2(addPrefixSpace bool, trimOffsets bool) *tokenizer.Tokenizer {
+	var model, err = bpe.NewBPEFromFS(fs, vocabFilename, mergeFilename)
 	if err != nil {
 		log.Fatal(err)
 	}