fix #2180

18a879f4 · Lysandre · d8034092 · 18a879f4
Commit 18a879f4 authored Dec 16, 2019 by Lysandre
Hide whitespace changes
Inline Side-by-side

Showing with 5 additions and 1 deletion

examples/run_generation.py examples/run_generation.py +5 -1

No files found.
--- a/examples/run_generation.py
+++ b/examples/run_generation.py
@@ -247,7 +247,11 @@ def main():
        out = out[:, len(context_tokens):].tolist()
        for o in out:
            text = tokenizer.decode(o, clean_up_tokenization_spaces=True)
-            text = text[: text.find(args.stop_token) if args.stop_token else None]
+            if args.stop_token:
+                index =  text.find(args.stop_token)
+                if index == -1:
+                    index = None
+                text = text[:index]
            print(text)