From 091037ec20743ac6c7bccb75b59743045a692c4a Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 2 Apr 2026 16:13:45 -0300
Subject: [PATCH] Fix top_logprobs_ids missing for llama.cpp loader

---
 modules/api/completions.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/modules/api/completions.py b/modules/api/completions.py
index 4eb8fdad..98bcff47 100644
--- a/modules/api/completions.py
+++ b/modules/api/completions.py
@@ -299,8 +299,9 @@ def format_completion_logprobs(entries):
             t = item.get('token', '')
             lp = item.get('logprob', item.get('prob', 0))
             top_dict[t] = lp
-            if 'token_id' in item:
-                top_dict_ids[item['token_id']] = lp
+            tid = item.get('token_id', item.get('id'))
+            if tid is not None:
+                top_dict_ids[tid] = lp
         top_logprobs.append(top_dict)
         top_logprobs_ids.append(top_dict_ids if top_dict_ids else None)