CI: Fix unittest for engine input token ids and output token ids (#2646)

This commit is contained in:
Shi Shuai
2024-12-29 21:28:59 +00:00
committed by GitHub
parent 35bdb48557
commit fad29f7f52
2 changed files with 18 additions and 31 deletions

View File

@@ -361,9 +361,13 @@ class BatchStrOut:
output_ids: Optional[List[int]]
# Token counts
# real input and output tokens can be get from
# origin_input_ids and output_ids by enabling --return_token_ids
# TODO (Shuai): Rename this to clarify the meaning.
prompt_tokens: List[int]
completion_tokens: List[int]
cached_tokens: List[int]
# Logprobs
input_token_logprobs_val: List[float]
input_token_logprobs_idx: List[int]