diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c34e2729..61899d2e2 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: retain end-of-generation token logprobs by @Kumar-Abhimanyu in #2383 + ## [0.3.36] - feat: update llama.cpp to ggml-org/llama.cpp@0c1e57098 diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 14e2f8500..6e0e59484 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1305,6 +1305,7 @@ def _create_completion( ) text: bytes = b"" returned_tokens: int = 0 + eog_logprobs: Optional[CompletionLogprobs] = None stop = ( stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else [] ) @@ -1411,6 +1412,34 @@ def logit_bias_processor( ): if llama_cpp.llama_vocab_is_eog(self._model.vocab, token): text = self.detokenize(completion_tokens, prev_tokens=prompt_tokens) + if logprobs is not None: + token_offset = len(prompt_tokens) + len(completion_tokens) - 1 + probabilities = Llama.logits_to_logprobs( + self._scores[token_offset, :] + ).tolist() + token_str = self.detokenize([token], special=True).decode( + "utf-8", errors="ignore" + ) + top_tokens = sorted( + range(len(probabilities)), + key=probabilities.__getitem__, + reverse=True, + )[:logprobs] + top_logprob = { + self.detokenize([i], special=True).decode( + "utf-8", errors="ignore" + ): probabilities[i] + for i in top_tokens + } + top_logprob[token_str] = probabilities[token] + eog_logprobs = { + "tokens": [token_str], + "text_offset": [ + len(prompt) + len(text.decode("utf-8", errors="ignore")) + ], + "token_logprobs": [probabilities[token]], + "top_logprobs": [top_logprob], + } finish_reason = "stop" break @@ -1699,7 +1728,7 @@ def logit_bias_processor( { "text": "", "index": 0, - "logprobs": None, + "logprobs": eog_logprobs, "finish_reason": finish_reason, } ], @@ -1791,6 +1820,11 @@ def logit_bias_processor( "token_logprobs": token_logprobs, "top_logprobs": top_logprobs, } + if eog_logprobs is not None: + tokens.extend(eog_logprobs["tokens"]) + text_offsets.extend(eog_logprobs["text_offset"]) + token_logprobs.extend(eog_logprobs["token_logprobs"]) + top_logprobs.extend(eog_logprobs["top_logprobs"]) yield { "id": completion_id, @@ -1851,6 +1885,9 @@ def create_completion( min_p: The min-p value to use for minimum p sampling. Minimum P sampling as described in https://github.com/ggerganov/llama.cpp/pull/3841 typical_p: The typical-p value to use for sampling. Locally Typical Sampling implementation described in the paper https://arxiv.org/abs/2202.00666. logprobs: The number of logprobs to return. If None, no logprobs are returned. + Includes the sampled end-of-generation token in probability metadata, + without adding it to completion text or usage. When streaming, its + logprobs are returned in the final chunk. echo: Whether to echo the prompt. stop: A list of strings to stop generation when encountered. frequency_penalty: The penalty to apply to tokens based on their frequency in the prompt. diff --git a/tests/test_llama.py b/tests/test_llama.py index c1c16e30c..c80d25676 100644 --- a/tests/test_llama.py +++ b/tests/test_llama.py @@ -16,6 +16,71 @@ MODEL = "./vendor/llama.cpp/models/ggml-vocab-llama-spm.gguf" +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("echo", [False, True]) +@pytest.mark.parametrize("prefix", [[], [15043]]) +@pytest.mark.parametrize("requested_logprobs", [None, 0, 3]) +def test_completion_eos_logprobs(monkeypatch, stream, echo, prefix, requested_logprobs): + model = llama_cpp.Llama( + model_path=MODEL, + vocab_only=True, + verbose=False, + logits_all=True, + n_ctx=32, + ) + eos = model.token_eos() + expected_logits = np.arange(model.n_vocab(), dtype=np.float32) / 1000 + + def generate(tokens, **kwargs): + model.n_tokens = len(tokens) + model.scores[: model.n_tokens] = 0 + for token in prefix + [eos]: + model.scores[model.n_tokens - 1] = expected_logits if token == eos else 0 + yield token + model.n_tokens += 1 + + monkeypatch.setattr(model, "generate", generate) + try: + output = model.create_completion( + "Hi", + max_tokens=8, + stream=stream, + echo=echo, + logprobs=requested_logprobs, + ) + # Streaming completions currently do not echo the prompt. + expected_text = ("Hi" if echo and not stream else "") + model.detokenize( + prefix + ).decode() + if stream: + chunks = list(output) + assert "".join(c["choices"][0]["text"] for c in chunks) == expected_text + final = chunks[-1]["choices"][0] + else: + final = output["choices"][0] + assert final["text"] == expected_text + assert output["usage"]["completion_tokens"] == len(prefix) + assert final["finish_reason"] == "stop" + probabilities = final["logprobs"] + if requested_logprobs is None: + assert probabilities is None + else: + eos_text = model.detokenize([eos], special=True).decode() + expected_probability = log_softmax(expected_logits)[eos] + assert probabilities["tokens"][-1] == eos_text + assert probabilities["token_logprobs"][-1] == pytest.approx( + expected_probability + ) + assert probabilities["top_logprobs"][-1][eos_text] == pytest.approx( + expected_probability + ) + assert probabilities["text_offset"][-1] == len("Hi") + len( + model.detokenize(prefix).decode() + ) + finally: + model.close() + + def test_llama_cpp_version(): assert llama_cpp.__version__