Fix skip special tokens in lmi-dist (deepjavalibrary#908)

KexinFeng · Aug 16, 2023 · c3fd994 · c3fd994
1 parent 0865639
commit c3fd994
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py b/engines/python/setup/djl_python/rolling_batch/lmi_dist_rolling_batch.py
@@ -104,7 +104,9 @@ def _prefill_and_decode(self, new_batch):
             is_last_token = generation.generated_text is not None
             if not is_last_token:
                 req_ids.append(r.id)
-            r.set_next_token(generation.token_text, last_token=is_last_token)
+            r.set_next_token(
+                "" if generation.token_is_special else generation.token_text,
+                last_token=is_last_token)
 
         # filter the requests that are stopped.
         if self.cache: