fix #2180

2019-12-16 16:44:29 -05:00
parent d803409215
commit 18a879f475
1 changed files with 5 additions and 1 deletions
--- 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)