Skip to content

Commit

Permalink
Update text_generation.py
Browse files Browse the repository at this point in the history
  • Loading branch information
TolyaTalamanov authored Jan 3, 2025
1 parent 64f97a1 commit dcbf890
Showing 1 changed file with 2 additions and 2 deletions.
4 changes: 2 additions & 2 deletions tools/llm_bench/task/text_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -225,7 +225,7 @@ def run_text_generation_genai(input_text, num, model, tokenizer, args, iter_data
log.info(out_str)
gen_config = model.get_generation_config()
gen_config.max_new_tokens = max_gen_tokens
gen_config.rng_seed= args["seed"]
gen_config.rng_seed = args["seed"]
gen_config.num_beams = args["num_beams"]
if args.get('draft_model', ''):
config_info = "Speculative decoding config: "
Expand Down Expand Up @@ -377,7 +377,7 @@ def run_text_generation_genai_with_stream(input_text, num, model, tokenizer, arg
max_gen_tokens = DEFAULT_OUTPUT_TOKEN_SIZE if args['infer_count'] is None else args['infer_count']
streamer.reset()
gen_config = model.get_generation_config()
gen_config.rng_seed= args["seed"]
gen_config.rng_seed = args["seed"]
gen_config.max_new_tokens = max_gen_tokens
gen_config.num_beams = args["num_beams"]
if args.get('draft_model', ''):
Expand Down

0 comments on commit dcbf890

Please sign in to comment.