Replace logits_to_logprobs implementation with numpy equivalent to llama.cpp (abetlen#991)

player1537 · web-flow · commit ef22e478db53 · 2023-12-11T20:46:27.000-05:00
See abetlen#990. This change makes the logits_to_logprobs function equivalent to the version in the llama.cpp repository. It uses numpy so it's much faster than the previous version.
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -2280,10 +2280,14 @@ def token_nl(self) -> int:
         return self._model.token_nl()
 
     @staticmethod
-    def logits_to_logprobs(logits: List[float]) -> List[float]:
-        exps = [math.exp(float(x)) for x in logits]
-        sum_exps = sum(exps)
-        return [math.log(x / sum_exps) for x in exps]
+    def logits_to_logprobs(logits: npt.NDArray[np.single]) -> npt.NDArray[np.single]:
+        maximum = np.max(logits)
+        tmp = np.subtract(logits, maximum, dtype=np.single)
+        np.exp(tmp, out=tmp)
+        normalizer = 1.0 / np.sum(tmp)
+        np.multiply(normalizer, tmp, out=tmp)
+        np.log(tmp, out=tmp)
+        return tmp
 
     @staticmethod
     def longest_token_prefix(a: Sequence[int], b: Sequence[int]):
