Get logits as numpy array

pygptj/model.py

@@ -18,6 +18,7 @@
 import sys
 import _pygptj as pp
 from pygptj._logger import set_log_level
+import numpy as np
 
 
 class Model:
@@ -36,6 +37,7 @@ def new_text_callback(text):
     ```
     """
     _new_text_callback = None
+    _logits_callback = None
 
     def __init__(self,
                  model_path: str,
@@ -64,6 +66,8 @@ def __init__(self,
 
         self.res = ""
 
+        self.logits = []
+
     def _load_model(self):
         """
         Helper function to load the model
@@ -84,10 +88,24 @@ def _call_new_text_callback(self, text_bytes) -> None:
             except UnicodeDecodeError:
                 logging.warning(f"UnicodeDecodeError of bytes {text_bytes}")
         # save res
+
+    def _call_logits_callback(self, logits: np.ndarray):
+        """
+        Internal logits_callback that saves the logit representation at each token.
+        :return: None
+        """
+        self.logits.append(logits.tolist())
+
+        if Model._logits_callback is not None:
+            Model._logits_callback(logits)
+
+    def braindump(self, path: str):
+        np.save(path, np.asarray(self.logits))
 
     def generate(self,
                  prompt: str,
                  new_text_callback: Callable[[str], None] = None,
+                 logits_callback: Callable = None,
                  n_predict: int = 128,
                  seed: int = -1,
                  n_threads: int = 4,
@@ -124,8 +142,11 @@ def generate(self,
         self.res = ""
         Model._new_text_callback = new_text_callback
 
+        # assign _logits_callback used for saving logits, token by token
+        Model._logits_callback = logits_callback
+
         # run the prediction
-        pp.gptj_generate(self.gpt_params, self._model, self._vocab, self._call_new_text_callback)
+        pp.gptj_generate(self.gpt_params, self._model, self._vocab, self._call_new_text_callback, self._call_logits_callback)
         return self.res
 
     @staticmethod

src/gptj.cpp

@@ -21,6 +21,7 @@
 #include <iostream>
 
 #include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
 namespace py = pybind11;
 
 
@@ -595,7 +596,7 @@ bool gptj_eval(
     return true;
 }
 
-int gptj_generate(gpt_params params, struct gptj_model & model, struct gpt_vocab & vocab,  py::function new_text_callback) {
+int gptj_generate(gpt_params params, struct gptj_model & model, struct gpt_vocab & vocab,  py::function new_text_callback, py::function logits_callback) {
 //    auto model = context->model;
 //    auto vocab = context->vocab;
 
@@ -656,7 +657,9 @@ int gptj_generate(gpt_params params, struct gptj_model & model, struct gpt_vocab
                 printf("Failed to predict\n");
                 return 1;
             }
-
+            // collect logits for each token
+            py::array_t<float> _logits = py::array_t<float>{model.hparams.n_vocab, logits.data(), py::none()};
+            logits_callback(_logits);
             t_predict_us += ggml_time_us() - t_start_us;
         }
 

src/gptj.h

@@ -85,4 +85,4 @@ bool gptj_eval(
               std::vector<float>         & embd_w,
               size_t                     & mem_per_token);
 
-int gptj_generate(gpt_params params, struct gptj_model & model, struct gpt_vocab & vocab, py::function new_text_callback);
+int gptj_generate(gpt_params params, struct gptj_model & model, struct gpt_vocab & vocab, py::function new_text_callback, py::function logits_callback);