diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..dff96bde9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +- fix: resize detokenization buffer for long token pieces by @grinder81 in #2388 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/llama_cpp/_internals.py b/llama_cpp/_internals.py index b45d34b2c..9d00500bb 100644 --- a/llama_cpp/_internals.py +++ b/llama_cpp/_internals.py @@ -196,7 +196,14 @@ def detokenize(self, tokens: List[int], special: bool = False) -> bytes: n = llama_cpp.llama_token_to_piece( self.vocab, llama_cpp.llama_token(token), buffer, size, 0, special ) - assert n <= size + if n < 0: + size = -n + buffer = (ctypes.c_char * size)() + n = llama_cpp.llama_token_to_piece( + self.vocab, llama_cpp.llama_token(token), buffer, size, 0, special + ) + if n < 0 or n > size: + raise RuntimeError("Token decoding failed after resizing buffer") output += bytes(buffer[:n]) # NOTE: Llama1 models automatically added a space at the start of the prompt # this line removes a leading space if the first token is a beginning of sentence token diff --git a/tests/test_detokenize_buffer.py b/tests/test_detokenize_buffer.py new file mode 100644 index 000000000..df0296557 --- /dev/null +++ b/tests/test_detokenize_buffer.py @@ -0,0 +1,58 @@ +import ctypes +from types import SimpleNamespace + +import pytest + +from llama_cpp import _internals as internals + + +@pytest.mark.parametrize("length", [0, 1, 31, 32, 33, 34, 128]) +def test_detokenize_resizes_for_required_length(monkeypatch, length): + payload = b"x" * length + sizes = [] + + def piece(vocab, token, buffer, size, lstrip, special): + sizes.append(size) + if size < len(payload): + return -len(payload) + ctypes.memmove(buffer, payload, len(payload)) + return len(payload) + + monkeypatch.setattr(internals.llama_cpp, "llama_token_to_piece", piece) + model = SimpleNamespace(vocab=None, token_bos=lambda: -1) + assert internals.LlamaModel.detokenize(model, [42]) == payload + assert sizes == ([32] if length <= 32 else [32, length]) + + +def test_detokenize_preserves_unicode_and_reuses_buffer(monkeypatch): + payloads = ["বাংলা ভাষার পরীক্ষা".encode(), "১২৩৪".encode()] + assert len(payloads[0]) > 32 + sizes = [] + + def piece(vocab, token, buffer, size, lstrip, special): + assert special is True + assert lstrip == 0 + payload = payloads[token.value] + sizes.append(size) + if size < len(payload): + return -len(payload) + ctypes.memmove(buffer, payload, len(payload)) + return len(payload) + + monkeypatch.setattr(internals.llama_cpp, "llama_token_to_piece", piece) + model = SimpleNamespace(vocab=None, token_bos=lambda: -1) + assert internals.LlamaModel.detokenize(model, [0, 1], special=True) == b"".join( + payloads + ) + assert sizes == [32, len(payloads[0]), len(payloads[0])] + + +@pytest.mark.parametrize("retry_result", [-40, 41]) +def test_detokenize_rejects_invalid_retry(monkeypatch, retry_result): + responses = iter([-40, retry_result]) + monkeypatch.setattr( + internals.llama_cpp, "llama_token_to_piece", lambda *args: next(responses) + ) + model = SimpleNamespace(vocab=None, token_bos=lambda: -1) + with pytest.raises(RuntimeError, match="after resizing buffer"): + internals.LlamaModel.detokenize(model, [42])