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:]

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)