chore(trtllm): minor refactoring

This commit is contained in:
Morgan Funtowicz 2024-10-21 23:40:20 +02:00
parent 60a08a283d
commit 848b8ad554

View File

@ -100,11 +100,10 @@ huggingface::tgi::backends::TensorRtLlmBackend::TensorRtLlmBackend(
SPDLOG_INFO(FMT_STRING("Found {:d} EOS tokens"), eosTokenIds.size()); SPDLOG_INFO(FMT_STRING("Found {:d} EOS tokens"), eosTokenIds.size());
stopWords = std::list<decltype(stopWords)::value_type>(eosTokenIds.size()); stopWords = std::list<decltype(stopWords)::value_type>(eosTokenIds.size());
std::transform(eosTokenIds.cbegin(), eosTokenIds.cend(), stopWords.begin(), const auto to_single_token = [](const auto tokenIdObj) -> decltype(stopWords)::value_type {
[](const auto tokenIdObj) -> decltype(stopWords)::value_type { return {tokenIdObj.template get<tle::TokenIdType>()};
const auto tokenId = tokenIdObj.template get<tle::TokenIdType>(); };
return {tokenId}; std::transform(eosTokenIds.cbegin(), eosTokenIds.cend(), stopWords.begin(), to_single_token);
});
} }
} else { } else {
SPDLOG_INFO("No EOS tokens found, generation_config.json doesn't exist"); SPDLOG_INFO("No EOS tokens found, generation_config.json doesn't exist");
@ -114,13 +113,13 @@ huggingface::tgi::backends::TensorRtLlmBackend::TensorRtLlmBackend(
[[nodiscard("Returned number of requests needs to be consumed")]] [[nodiscard("Returned number of requests needs to be consumed")]]
size_t huggingface::tgi::backends::TensorRtLlmBackend::NumResponsesReady() const { size_t huggingface::tgi::backends::TensorRtLlmBackend::NumResponsesReady() const {
#ifdef NDEBUG
return executor.getNumResponsesReady();
#else
const auto numResponses = executor.getNumResponsesReady(); const auto numResponses = executor.getNumResponsesReady();
#ifndef NDEBUG
if (numResponses > 0) SPDLOG_INFO(FMT_STRING("Num responses ready: {:d}"), numResponses); if (numResponses > 0) SPDLOG_INFO(FMT_STRING("Num responses ready: {:d}"), numResponses);
#endif
return numResponses; return numResponses;
#endif
} }
[[nodiscard("Returned request id needs to be provided back to gather generated tokens")]] [[nodiscard("Returned request id needs to be provided back to gather generated tokens")]]