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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user