Workaround: second-to-last token is always " ", remove this manually

LLaVA-VL · Oct 17, 2023 · f2e664f · f2e664f
1 parent 86efd49
commit f2e664f
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/predict.py b/predict.py
@@ -127,11 +127,23 @@ def predict(
                 use_cache=True,
                 stopping_criteria=[stopping_criteria]))
             thread.start()
+            # workaround: second-to-last token is always " "
+            # but we want to keep it if it's not the second-to-last token
+            prepend_space = False
             for new_text in streamer:
+                if new_text == " ":
+                    prepend_space = True
+                    continue
                 if new_text.endswith(stop_str):
                     new_text = new_text[:-len(stop_str)].strip()
+                    prepend_space = False
+                elif prepend_space:
+                    new_text = " " + new_text
+                    prepend_space = False
                 if len(new_text):
                     yield new_text
+            if prepend_space:
+                yield " "
             thread.join()