fix: resize detokenization buffer for long token pieces

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

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

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