remove LOGPROBS_TYPE

serve/mlc_serve/engine/base.py

@@ -9,6 +9,7 @@
 import numpy as np
 
 from .sampling_params import SamplingParams, SamplingType
+from ..api.protocol import LogprobsContent
 
 RequestId = str
 LOGPROBS_TYPE = Tuple[Tuple, List[Tuple]]
@@ -172,7 +173,7 @@ class SequenceOutput:
     finish_reason: Optional[FinishReason] = None
     # Number of generated tokens so far
     num_generated_tokens: int = 0
-    logprob_info: Optional[LOGPROBS_TYPE] = None
+    logprob_info: Optional[LogprobsContent] = None
 
     @property
     def is_finished(self) -> bool:

serve/mlc_serve/engine/engine_common.py

@@ -169,7 +169,7 @@ def logprob_detokenize(
             ))
 
     logprobs_content = LogprobsContent(
-        token=tokenizer.decode(logprob_info.current_token),
+        token=tokenizer.decode([logprob_info.current_token]),
         logprob=logprob_info,
         # TODO(vvchernov): implement bytes based on https://platform.openai.com/docs/api-reference/chat/object
         bytes=None,

serve/mlc_serve/engine/staging_engine_worker.py

@@ -6,14 +6,12 @@
 import multiprocessing.synchronize
 from dataclasses import dataclass
 from threading import Thread, Lock
-from typing import Callable, Optional, Union, Tuple, Any, Dict, List
+from typing import Callable, Optional, Union, Any, Dict, List
 
 import structlog
-import numpy as np
 
 from .base import (
     FinishReason,
-    LOGPROBS_TYPE,
     RequestId,
     RequestState,
     ValidationError,
@@ -32,6 +30,7 @@
     EngineBase,
 )
 from ..logging_utils import configure_logging
+from ..api.protocol import LogprobsContent
 
 LOG = structlog.stdlib.get_logger(__name__)
 
@@ -67,7 +66,7 @@ class SequenceGenerationOutput:
     new_tokens: List[int]
     finish_reason: Optional[FinishReason] = None
     error: Optional[Union[str, ValidationError]] = None
-    logprob_info: Optional[LOGPROBS_TYPE] = None
+    logprob_info: Optional[LogprobsContent] = None
 
 
 @dataclass