handle empty encode/decode params

daulet · May 7, 2023 · f3f1a0e · f3f1a0e
1 parent 02705a2
commit f3f1a0e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/tokenizer.go b/tokenizer.go
@@ -73,6 +73,9 @@ func (t *Tokenizer) Encode(str string, addSpecialTokens bool) []uint32 {
 }
 
 func (t *Tokenizer) Decode(tokenIDs []uint32, skipSpecialTokens bool) string {
+	if len(tokenIDs) == 0 {
+		return ""
+	}
 	len := C.uint(len(tokenIDs))
 	res := C.decode(t.tokenizer, (*C.uint)(unsafe.Pointer(&tokenIDs[0])), len, C.bool(skipSpecialTokens))
 	defer C.free(unsafe.Pointer(res))

diff --git a/tokenizer_test.go b/tokenizer_test.go
@@ -76,6 +76,18 @@ func TestEncode(t *testing.T) {
 			addSpecial: true,
 			want:       []uint32{101, 2829, 4419, 14523, 2058, 1996, 13971, 3899, 102},
 		},
+		{
+			name:       "empty string",
+			str:        "",
+			addSpecial: false,
+			want:       []uint32{},
+		},
+		{
+			name:       "empty string with special tokens",
+			str:        "",
+			addSpecial: false,
+			want:       []uint32{},
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
@@ -174,6 +186,12 @@ func TestDecode(t *testing.T) {
 			skipSpecial: false,
 			want:        "[CLS] brown fox jumps over the lazy dog [SEP]",
 		},
+		{
+			name:        "no tokens",
+			tokens:      []uint32{},
+			skipSpecial: false,
+			want:        "",
+		},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {