Colour tokens based on how perplexing they are

Author: blackle

llama.cpp

@@ -1019,7 +1019,8 @@ static llama_vocab::id llama_sample_top_p_top_k(
         int top_k,
         double top_p,
         double temp,
-        double repeat_penalty) {
+        double repeat_penalty,
+        double* out_prob) {
     auto & rng = lctx.rng;
 
     const auto & vocab = lctx.vocab;
@@ -1098,6 +1099,7 @@ static llama_vocab::id llama_sample_top_p_top_k(
     std::discrete_distribution<> dist(probs.begin(), probs.end());
     int idx = dist(rng);
 
+    *out_prob = probs[idx];
     return logits_id[idx].second;
 }
 
@@ -1505,7 +1507,8 @@ llama_token llama_sample_top_p_top_k(
                     int   top_k,
                  double   top_p,
                  double   temp,
-                 double   repeat_penalty) {
+                 double   repeat_penalty,
+                 double*  out_prob) {
     const int64_t t_start_sample_us = ggml_time_us();
 
     llama_token result = 0;
@@ -1519,7 +1522,8 @@ llama_token llama_sample_top_p_top_k(
             top_k,
             top_p,
             temp,
-            repeat_penalty);
+            repeat_penalty,
+            out_prob);
 
     ctx->t_sample_us += ggml_time_us() - t_start_sample_us;
     ctx->n_sample++;

llama.h

@@ -123,7 +123,8 @@ extern "C" {
                         int   top_k,
                      double   top_p,
                      double   temp,
-                     double   repeat_penalty);
+                     double   repeat_penalty,
+                     double*  out_prob);
 
     // Performance information
     LLAMA_API void llama_print_timings(struct llama_context * ctx);

main.cpp

@@ -154,6 +154,15 @@ void sigint_handler(int signo) {
 }
 #endif
 
+void print_color_for_prob(double prob) {
+    int p = prob*0xFF;
+    printf("\x1b[48;2;%d;%d;%dm", 0xFF-p, (0xFF-p)/10 + 10, p/10);
+}
+
+void print_color_reset() {
+    printf("\x1b[0m");
+}
+
 int main(int argc, char ** argv) {
     // has to be called once at the start of the program to init ggml stuff
     ggml_time_init();
@@ -335,6 +344,7 @@ int main(int argc, char ** argv) {
 
         n_past += embd.size();
         embd.clear();
+        double prob = 0.0;
 
         if ((int) embd_inp.size() <= input_consumed) {
             // out of user input, sample next token
@@ -356,7 +366,7 @@ int main(int argc, char ** argv) {
                     logits[llama_token_eos()] = 0;
                 }
 
-                id = llama_sample_top_p_top_k(ctx, last_n_tokens.data(), last_n_tokens.size(), top_k, top_p, temp, repeat_penalty);
+                id = llama_sample_top_p_top_k(ctx, last_n_tokens.data(), last_n_tokens.size(), top_k, top_p, temp, repeat_penalty, &prob);
 
                 last_n_tokens.erase(last_n_tokens.begin());
                 last_n_tokens.push_back(id);
@@ -393,13 +403,15 @@ int main(int argc, char ** argv) {
             }
         }
 
+        print_color_for_prob(prob);
         // display text
         if (!input_noecho) {
             for (auto id : embd) {
                 printf("%s", llama_token_to_str(ctx, id));
             }
             fflush(stdout);
         }
+        print_color_reset();
         // reset color to default if we there is no pending user input
         if (!input_noecho && (int)embd_inp.size() == input_consumed) {
             set_console_state(CONSOLE_STATE_DEFAULT);