use decode instead batch decode for streaming (deepjavalibrary#1016)

david-sitsky · Aug 12, 2023 · 4bc70f7 · 4bc70f7
1 parent f5cc091
commit 4bc70f7
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/streaming_utils.py b/engines/python/setup/djl_python/streaming_utils.py
@@ -38,8 +38,12 @@ def __init__(self, tokenizer, **decode_kwargs):
 
     def put(self, value):
         self.started = True
-        text = self.tokenizer.batch_decode(value, **self.decode_kwargs)
-        self.queue.put(text)
+        if len(value.shape) == 1:
+            value = [value]
+        result = []
+        for item in value:
+            result.append(self.tokenizer.decode(item, **self.decode_kwargs))
+        self.queue.put(result)
 
     def put_text(self, value):
         self.queue.put(value)