file path support

pkoukk · Jun 27, 2023 · 0227988 · 0227988
1 parent 35a2ee5
commit 0227988
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 1 deletion.
diff --git a/loader_test.go b/loader_test.go
@@ -0,0 +1,16 @@
+package tiktoken_loader
+
+import (
+	"testing"
+)
+
+func TestOffline(t *testing.T) {
+	loader := NewOfflineLoader()
+	bpeRanks, err := loader.LoadTiktokenBpe("https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken")
+	if err != nil {
+		t.Error(err)
+	}
+	if len(bpeRanks) == 0 {
+		t.Error("bpeRanks is empty")
+	}
+}
diff --git a/offline_loader.go b/offline_loader.go
@@ -2,6 +2,7 @@ package tiktoken_loader
 
 import (
 	"encoding/base64"
+	"path"
 	"strconv"
 	"strings"
 
@@ -11,7 +12,8 @@ import (
 type OfflineLoader struct{}
 
 func (l *OfflineLoader) LoadTiktokenBpe(tiktokenBpeFile string) (map[string]int, error) {
-	contents, err := assets.Assets.ReadFile(tiktokenBpeFile)
+	baseFileName := path.Base(tiktokenBpeFile)
+	contents, err := assets.Assets.ReadFile(baseFileName)
 	if err != nil {
 		return nil, err
 	}