fix(function_call): group batch decode by options instead of fallback (#16698)

Co-authored-by: Muqi Li <muqi1029@gmail.com>
This commit is contained in:
若可
2026-01-10 11:38:13 +08:00
committed by GitHub
parent 9c64a15ad4
commit fbc128a32e

View File

@@ -17,8 +17,8 @@ import dataclasses
import logging
import os
import signal
from collections import OrderedDict
from typing import Dict, List, Union
from collections import OrderedDict, defaultdict
from typing import Dict, List, Tuple, Union
import psutil
import pybase64
@@ -174,6 +174,46 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
# If it is embedding model, no detokenization is needed.
return recv_obj
def _grouped_batch_decode(
self,
ids_list: List[List[int]],
skip_list: List[bool],
space_list: List[bool],
) -> List[str]:
"""Batch decode with grouping by (skip_special_tokens, spaces_between_special_tokens)."""
assert self.tokenizer is not None
# fast path
first_skip, first_space = skip_list[0], space_list[0]
if all(s == first_skip for s in skip_list) and all(
sp == first_space for sp in space_list
):
return self.tokenizer.batch_decode(
ids_list,
skip_special_tokens=first_skip,
spaces_between_special_tokens=first_space,
)
# Group indices by (skip, space) tuple
groups: Dict[Tuple[bool, bool], List[int]]
groups = defaultdict(list)
for idx, (skip, space) in enumerate(zip(skip_list, space_list)):
groups[(skip, space)].append(idx)
# Decode each group and collect results
results: List[str] = [""] * len(ids_list)
for (skip, space), indices in groups.items():
decoded = self.tokenizer.batch_decode(
[ids_list[idx] for idx in indices],
skip_special_tokens=skip,
spaces_between_special_tokens=space,
)
for idx, text in zip(indices, decoded):
results[idx] = text
return results
def _decode_batch_token_id_output(self, recv_obj: BatchTokenIDOutput):
bs = len(recv_obj.rids)
@@ -203,27 +243,18 @@ class DetokenizerManager(MultiHttpWorkerDetokenizerMixin):
surr_ids.append(s.decode_ids[s.surr_offset : s.read_offset])
# Decode token ids to strings
# TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
skip_uniform = len(set(recv_obj.skip_special_tokens)) == 1
space_uniform = len(set(recv_obj.spaces_between_special_tokens)) == 1
if not self.disable_tokenizer_batch_decode or not (
skip_uniform and space_uniform
):
if not self.disable_tokenizer_batch_decode:
if not self.is_dummy:
# Run normal batch decode
surr_texts = self.tokenizer.batch_decode(
surr_texts = self._grouped_batch_decode(
surr_ids,
skip_special_tokens=recv_obj.skip_special_tokens[0],
spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
0
],
recv_obj.skip_special_tokens,
recv_obj.spaces_between_special_tokens,
)
read_texts = self.tokenizer.batch_decode(
read_texts = self._grouped_batch_decode(
read_ids,
skip_special_tokens=recv_obj.skip_special_tokens[0],
spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
0
],
recv_obj.skip_special_tokens,
recv_obj.spaces_between_special_tokens,
)
else:
# If it is dummy weights, just return dummy strings to prevent potential detokenization edge cases