From 994fa52c0b99df1656b609a14eb431c80292dbcb Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Tue, 10 Dec 2024 23:28:13 +0100
Subject: [PATCH] feat(trtllm): fix logits retrieval

---
 backends/trtllm/csrc/ffi.hpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/backends/trtllm/csrc/ffi.hpp b/backends/trtllm/csrc/ffi.hpp
index d14578b8..840614bb 100644
--- a/backends/trtllm/csrc/ffi.hpp
+++ b/backends/trtllm/csrc/ffi.hpp
@@ -51,10 +51,11 @@ namespace huggingface::tgi::backends::trtllm {
         const auto reqId = r.getRequestId();
         if (!r.hasError()) [[likely]] {
             const auto result = r.getResult();
+            const auto logits = result.logProbs.value()[0];
             return generation_step_t{
                     reqId,
                     static_cast<uint32_t>(result.outputTokenIds[0][0]),
-                    result.logProbs.value()[0][0],
+                    logits.back(),
                     result.isFinal,
                     as_finish_reason_t(result.finishReasons[0]),
                     false,