skip_prompt=True in TextIteratorStreamer (#2193)

* skip_prompt in `TextIteratorStreamer` * Update trl/commands/cli.py Co-authored-by: Kashif Rasul <[email protected]> * Update generation streamer in chat.py --------- Co-authored-by: Kashif Rasul <[email protected]>
huggingface · Oct 7, 2024 · adf58d8 · adf58d8
1 parent 9aa0225
commit adf58d8
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/examples/scripts/chat.py b/examples/scripts/chat.py
@@ -273,7 +273,7 @@ def chat_cli():
         user = args.user
 
     model, tokenizer = load_model_and_tokenizer(args)
-    generation_streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    generation_streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True, skip_prompt=True)
 
     pad_token_id, eos_token_ids = parse_eos_tokens(tokenizer, args.eos_tokens, args.eos_token_ids)
 

diff --git a/trl/commands/cli.py b/trl/commands/cli.py
@@ -96,6 +96,7 @@ def train(command_name):
             encoding="utf-8",
             cwd=os.getcwd(),
             env=os.environ.copy(),
+            capture_output=True,
         )
     except (CalledProcessError, ChildProcessError) as exc:
         console.log(f"TRL - {command_name.upper()} failed on ! See the logs above for further details.")