Don't convert logprobs arrays to lists #1021

kddubey · 2023-12-17T15:44:57Z

Follow-up to #1002

Speedup

Run this in an ipython shell (python -m pip install jupyter if you don't have it already)

import numpy as np
from llama_cpp import Llama

_scores: np.ndarray = -np.random.default_rng(123).uniform(
    low=0, high=60, size=(20, 32_000)
)
token_offset = 1

print("time old")
%timeit [Llama.logits_to_logprobs(row.tolist()) for row in _scores][token_offset:]
print()
print("time new")
%timeit Llama.logits_to_logprobs(_scores)[token_offset:]

prints

time old
72.6 ms ± 862 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

time new
24.5 ms ± 142 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

How has this been tested?

import numpy as np
from llama_cpp import Llama

_scores: np.ndarray = -np.random.default_rng(123).uniform(
    low=0, high=60, size=(20, 32_000)
)
token_offset = 1

# 1-D case
logits = _scores[token_offset, :]
current_logprobs_old = Llama.logits_to_logprobs(logits.tolist())
current_logprobs_new = Llama.logits_to_logprobs(logits)
assert np.allclose(current_logprobs_old, current_logprobs_new)

# 2-D case
all_logprobs_old = [Llama.logits_to_logprobs(row.tolist()) for row in _scores][
    token_offset:
]
all_logprobs_new = Llama.logits_to_logprobs(_scores)[token_offset:]
assert np.allclose(all_logprobs_old, all_logprobs_new)

kddubey · 2023-12-17T15:46:14Z

llama_cpp/llama.py

-            all_logprobs = [
-                Llama.logits_to_logprobs(row.tolist()) for row in self._scores
-            ][token_offset:]
+            all_logprobs = Llama.logits_to_logprobs(self._scores)[token_offset:]


all_logprobs is now a 2-D array instead of a list of 1-D arrays. But this change should be ok b/c it gets zipped, which will iterate over the first axis

abetlen · 2023-12-18T19:28:28Z

@kddubey thank you!

kddubey commented Dec 17, 2023

View reviewed changes

kddubey changed the title ~~Don't convert logprobs lists to arrays~~ Don't convert logprobs arrays to lists Dec 17, 2023

Don't convert logprobs arrays to lists

375d86b

abetlen merged commit 6b2e0e0 into abetlen:main Dec 18, 2023

kddubey deleted the lists-to-arrays branch December 19, 2023 03:23

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Don't convert logprobs arrays to lists #1021

Don't convert logprobs arrays to lists #1021

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Don't convert logprobs arrays to lists #1021

Don't convert logprobs arrays to lists #1021

Uh oh!

Conversation

Speedup

How has this been tested?

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!