fix: retain end-of-generation token logprobs

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

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.

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__