huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-22 15:32:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(trtllm): fix logits retrieval

Browse Source
This commit is contained in:
Morgan Funtowicz 2024-12-10 23:28:13 +01:00
parent 60059b6968
commit b653605e54
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
backends/trtllm/csrc/ffi.hpp
View File

@ -51,10 +51,11 @@ namespace huggingface::tgi::backends::trtllm {
const auto reqId = r.getRequestId(); const auto reqId = r.getRequestId();
if (!r.hasError()) [[likely]] { if (!r.hasError()) [[likely]] {
const auto result = r.getResult(); const auto result = r.getResult();
const auto logits = result.logProbs.value()[0];
return generation_step_t{ return generation_step_t{
reqId, reqId,
static_cast<uint32_t>(result.outputTokenIds[0][0]), static_cast<uint32_t>(result.outputTokenIds[0][0]),
result.logProbs.value()[0][0], logits.back(),
result.isFinal, result.isFinal,
as_finish_reason_t(result.finishReasons[0]), as_finish_reason_t(result.finishReasons[0]),
false, false,