Add skip_tokenizer_init args. (#959)

Co-authored-by: lzhang <zhanglei@modelbest.cn>
This commit is contained in:
gryffindor-rr
2024-08-10 03:14:13 +08:00
committed by GitHub
parent b16e856f11
commit 9cf0a5bada
10 changed files with 218 additions and 71 deletions

View File

@@ -59,11 +59,14 @@ class DetokenizerManager:
self.send_to_tokenizer = context.socket(zmq.PUSH)
self.send_to_tokenizer.connect(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
self.tokenizer = get_tokenizer(
server_args.tokenizer_path,
tokenizer_mode=server_args.tokenizer_mode,
trust_remote_code=server_args.trust_remote_code,
)
if server_args.skip_tokenizer_init:
self.tokenizer = None
else:
self.tokenizer = get_tokenizer(
server_args.tokenizer_path,
tokenizer_mode=server_args.tokenizer_mode,
trust_remote_code=server_args.trust_remote_code,
)
self.decode_status = {}
@@ -85,6 +88,11 @@ class DetokenizerManager:
assert isinstance(recv_obj, BatchTokenIDOut)
bs = len(recv_obj.rids)
if self.tokenizer is None:
# Send BatchTokenIDOut if no tokenizer init'ed.
self.send_to_tokenizer.send_pyobj(recv_obj)
continue
# Initialize decode status
read_ids, surr_ids = [], []
for i in range(bs):