From 2a20f48efad692a8c2744f10c673bbdbe0c751b7 Mon Sep 17 00:00:00 2001
From: Val Kharitonov <mail@kharvd.com>
Date: Mon, 13 Mar 2023 12:24:18 -0400
Subject: [PATCH] Fix UTF-8 handling (including colors) (#79)

---
 convert-pth-to-ggml.py | 29 ++++++++++++++++++++++-------
 main.cpp               |  9 +++++----
 2 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/convert-pth-to-ggml.py b/convert-pth-to-ggml.py
index fc217c7..d255750 100644
--- a/convert-pth-to-ggml.py
+++ b/convert-pth-to-ggml.py
@@ -22,7 +22,6 @@ import json
 import struct
 import numpy as np
 import torch
-
 from sentencepiece import SentencePieceProcessor
 
 if len(sys.argv) < 3:
@@ -101,12 +100,28 @@ for p in range(n_parts):
 
     # Is this correct??
     for i in range(32000):
-        # TODO: this is probably wrong - not sure how this tokenizer works
-        text = tokenizer.decode([29889, i]).encode('utf-8')
-        # remove the first byte (it's always '.')
-        text = text[1:]
-        fout.write(struct.pack("i", len(text)))
-        fout.write(text)
+        if tokenizer.is_unknown(i):
+            # "<unk>" token (translated as ??)
+            text = " \u2047 ".encode("utf-8")
+            fout.write(struct.pack("i", len(text)))
+            fout.write(text)
+        elif tokenizer.is_control(i):
+            # "<s>"/"</s>" tokens
+            fout.write(struct.pack("i", 0))
+        elif tokenizer.is_byte(i):
+            # "<U+XX>" tokens (which may be invalid UTF-8)
+            piece = tokenizer.id_to_piece(i)
+            if len(piece) != 6:
+                print("Invalid token: " + piece)
+                sys.exit(1)
+            byte_value = int(piece[3:-1], 16)
+            fout.write(struct.pack("i", 1))
+            fout.write(struct.pack("B", byte_value))
+        else:
+            # normal token. Uses U+2581 (LOWER ONE EIGHTH BLOCK) to represent spaces.
+            text = tokenizer.id_to_piece(i).replace("\u2581", " ").encode("utf-8")
+            fout.write(struct.pack("i", len(text)))
+            fout.write(text)
 
     for k, v in model.items():
         name = k
diff --git a/main.cpp b/main.cpp
index 98ccde5..d068761 100644
--- a/main.cpp
+++ b/main.cpp
@@ -939,6 +939,11 @@ int main(int argc, char ** argv) {
                     break;
                 }
             }
+
+            // reset color to default if we there is no pending user input
+            if (!input_noecho && params.use_color && embd_inp.size() == input_consumed) {
+                printf(ANSI_COLOR_RESET);
+            }
         }
 
         // display text
@@ -946,10 +951,6 @@ int main(int argc, char ** argv) {
             for (auto id : embd) {
                 printf("%s", vocab.id_to_token[id].c_str());
             }
-            // reset color to default if we there is no pending user input
-            if (params.use_color && embd_inp.size() <= input_consumed) {
-                printf(ANSI_COLOR_RESET);
-            }
             fflush(stdout);
         }