From 7b6ae8904178cb8f269719545a0396e5f726ff22 Mon Sep 17 00:00:00 2001
From: Georgi Gerganov <ggerganov@gmail.com>
Date: Thu, 17 Aug 2023 12:27:26 +0300
Subject: [PATCH] llama : fix tokenizer to use llama_char_to_byte

---
 llama.cpp                  | 22 +++++++++++++++++-----
 tests/test-tokenizer-0.cpp | 11 ++++++-----
 2 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/llama.cpp b/llama.cpp
index 76de7815f..fd3690432 100644
--- a/llama.cpp
+++ b/llama.cpp
@@ -2303,6 +2303,18 @@ static uint8_t llama_byte_to_char(const llama_vocab & vocab, uint8_t byte) {
     return false;
 }
 
+static uint8_t llama_char_to_byte(const llama_vocab & vocab, uint8_t ch) {
+    if (llama_vocab_type(vocab) == "spm") {
+        return ch + 3;
+    }
+
+    if (llama_vocab_type(vocab) == "bpe") {
+        return ch - 32;
+    }
+
+    return false;
+}
+
 static std::string llama_escape_whitespace(const std::string& text) {
     std::string result;
     bool escaping = false;
@@ -2439,7 +2451,7 @@ private:
         if (p == rev_merge.end()) {
             // output any symbols that did not form tokens as bytes.
             for (int j = 0; j < (int)symbol.n; ++j) {
-                llama_vocab::id token_id = llama_byte_to_char(vocab_, symbol.text[j]);
+                llama_vocab::id token_id = llama_char_to_byte(vocab_, symbol.text[j]);
                 output.push_back(token_id);
             }
             return;
@@ -4871,8 +4883,8 @@ int llama_token_to_str_with_model(const struct llama_model * model, llama_token
     return 0;
 }
 
-int llama_token_to_str(const struct llama_context * ctx, llama_token token, char * str, int length) {
-    return llama_token_to_str_with_model(&ctx->model, token, str, length);
+int llama_token_to_str(const struct llama_context * ctx, llama_token token, char * buf, int length) {
+    return llama_token_to_str_with_model(&ctx->model, token, buf, length);
 }
 
 std::string llama_token_to_str(const struct llama_context * ctx, llama_token token) {
@@ -4889,13 +4901,13 @@ std::string llama_token_to_str(const struct llama_context * ctx, llama_token tok
     return std::string(result.data(), result.size());
 }
 
-int llama_token_to_str_bpe(const struct llama_context * ctx, llama_token token, char * str, int length) {
+int llama_token_to_str_bpe(const struct llama_context * ctx, llama_token token, char * buf, int length) {
     if (0 <= token && token < llama_n_vocab_from_model(&ctx->model)) {
         std::string result = ctx->model.vocab.id_to_token[token].tok;
         if (length < (int) result.length()) {
             return -result.length();
         }
-        memcpy(str, result.c_str(), result.length());
+        memcpy(buf, result.c_str(), result.length());
         return result.length();
     }
     return 0;
diff --git a/tests/test-tokenizer-0.cpp b/tests/test-tokenizer-0.cpp
index f973271a3..9f5382679 100644
--- a/tests/test-tokenizer-0.cpp
+++ b/tests/test-tokenizer-0.cpp
@@ -89,6 +89,8 @@ int main(int argc, char **argv) {
         return 2;
     }
 
+    bool success = true;
+
     for (const auto & test_kv : k_tests()) {
         std::vector<llama_token> res = llama_tokenize(ctx, test_kv.first, true);
         fprintf(stderr, "%s : '%s' tokenized to '%s'\n",
@@ -103,7 +105,8 @@ int main(int argc, char **argv) {
         }
 
         if (!correct) {
-            fprintf(stderr, "%s : failed test: '%s'\n", __func__, test_kv.first.c_str());
+            fprintf(stderr, "%s : failed test:    '%s'\n", __func__, test_kv.first.c_str());
+            fprintf(stderr, "%s : detokenized to: '%s'\n", __func__, unescape_whitespace(ctx, test_kv.second).c_str());
             fprintf(stderr, "%s : expected tokens: ", __func__);
             for (const auto & t : test_kv.second) {
                 fprintf(stderr, "%6d, ", t);
@@ -115,9 +118,7 @@ int main(int argc, char **argv) {
             }
             fprintf(stderr, "\n");
 
-            llama_free_model(model);
-            llama_free(ctx);
-            return 3;
+            success = false;
         }
     }
 
@@ -126,5 +127,5 @@ int main(int argc, char **argv) {
 
     llama_backend_free();
 
-    return 0;
+    return success ? 0 : 3;
 }