From 0212b1774aea0b15ab246f728c3976100b7b7c93 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Wed, 17 Jul 2024 22:33:10 +0000 Subject: [PATCH] correctly forward back the log probabilities --- backends/trtllm/src/ffi.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backends/trtllm/src/ffi.cpp b/backends/trtllm/src/ffi.cpp index 39c7104c..f3b68da8 100644 --- a/backends/trtllm/src/ffi.cpp +++ b/backends/trtllm/src/ffi.cpp @@ -44,12 +44,12 @@ size_t huggingface::tgi::backends::TensorRtLlmBackendImpl::StreamTokens( const auto token = decoded.outputTokenIds[0][0]; const auto isFinal = decoded.isFinal; -// const auto logProb = decoded.logProbs.value()[0][0]; + const auto logProb = decoded.logProbs.value()[0][0]; ++numTokens; SPDLOG_DEBUG(FMT_STRING("\tStreamTokens -> {:d} {:.2f} (final = {})"), token, logProb, isFinal); - callback(std::move(ctx), token, 1.0, isFinal); + callback(std::move(ctx), token, logProb, isFinal); SPDLOG_DEBUG("\tStreamTokens -> Post callback"); } else { // TODO : Return rest::Result with error