feat(trtllm): fix logits retrieval

huggingface · Dec 10, 2024 · b653605 · b653605
1 parent 60059b6
commit b653605
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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,