From 6e0642ca195e2716e6e1bfa05d2b76f1e2dae0a3 Mon Sep 17 00:00:00 2001 From: Asghar Ghorbani Date: Tue, 4 Jun 2024 16:18:38 +0200 Subject: [PATCH] fix: fix logprobs when BOS is not present (#1471) * Fix lobprobs when BOS is not present * Fix logprobs when bos is not available --- llama_cpp/llama.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py index 45e1526..fb26950 100644 --- a/llama_cpp/llama.py +++ b/llama_cpp/llama.py @@ -1410,8 +1410,8 @@ class Llama: top_logprobs: List[Optional[Dict[str, float]]] = [] if echo: - # Remove leading BOS token - all_tokens = prompt_tokens[1:] + completion_tokens + # Remove leading BOS token if exists + all_tokens = prompt_tokens[1 if prompt_tokens[0] == self.token_bos() else 0:] + completion_tokens else: all_tokens = completion_tokens