fix mypy

octoml · Dec 29, 2023 · a822016 · a822016
1 parent a1047a2
commit a822016
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/serve/mlc_serve/engine/staging_engine.py b/serve/mlc_serve/engine/staging_engine.py
@@ -26,7 +26,7 @@
     update_sequence,
     logprob_detokenize
 )
-from .model_module import ModelModule, TokenizerModule, Tokenizer
+from .model_module import ModelModule, TokenizerModule
 from .staging_engine_worker import (
     AddRequestsCommand,
     CancelRequestCommand,
@@ -263,7 +263,7 @@ def step(self) -> InferenceStepResult:
                                 0,
                                 delta=delta,
                                 num_generated_tokens=(
-                                    len(state.token_ids) - state.prompt_len
+                                    len(state.prompt_token_ids) - state.prompt_len
                                 ),
                                 finish_reason=seq_output.finish_reason,
                                 logprob_info=logprob_detokenize(self.tokenizer, seq_output.logprob_info),

diff --git a/serve/mlc_serve/engine/sync_engine.py b/serve/mlc_serve/engine/sync_engine.py
@@ -233,7 +233,7 @@ def step(self) -> InferenceStepResult:
                             0,
                             delta=delta,
                             num_generated_tokens=(
-                                len(state.token_ids) - state.prompt_len
+                                len(state.prompt_token_ids) - state.prompt_len
                             ),
                             logprob_info=res.logprob_info
                         ),

diff --git a/serve/mlc_serve/model/paged_cache_model.py b/serve/mlc_serve/model/paged_cache_model.py
@@ -537,9 +537,9 @@ def generate(
             for i, (sequence_id, new_token) in enumerate(
                 zip(sequence_ids, next_tokens)
             ):
-                if not new_token in requests[i].sampling_params.appeared_tokens_freq:
-                    requests[i].sampling_params.appeared_tokens_freq[new_token] = 0
-                requests[i].sampling_params.appeared_tokens_freq[new_token] += 1
+                if not new_token in requests[index].sampling_params.appeared_tokens_freq:
+                    requests[index].sampling_params.appeared_tokens_freq[new_token] = 0
+                requests[index].sampling_params.appeared_tokens_freq[new_token] += 1
                 if sequence_id.sequence_index == PROMPT_SEQEUNCE_INDEX:
                     for seq_id in range(num_sequences[i]):
                         outputs.append(
@@ -581,9 +581,9 @@ def generate(
 
                 if maybe_new_token is not None:
                     new_token = maybe_new_token[0]
-                    if not new_token in requests[i].sampling_params.appeared_tokens_freq:
-                        requests[i].sampling_params.appeared_tokens_freq[new_token] = 0
-                    requests[i].sampling_params.appeared_tokens_freq[new_token] += 1
+                    if not new_token in requests[index].sampling_params.appeared_tokens_freq:
+                        requests[index].sampling_params.appeared_tokens_freq[new_token] = 0
+                    requests[index].sampling_params.appeared_tokens_freq[new_token] += 1
                     if sequence_id.sequence_index == PROMPT_SEQEUNCE_INDEX:
                         for seq_id in range(num_sequences[i]):
                             outputs.append(