Lucasoppem commited on
Commit
d72fd9e
·
verified ·
1 Parent(s): 2f9a5a7

Fix bug: Prevent inconsistent return value with `eos_early_stop`

Browse files

When `eos_early_stop = True`, generate() returns **prompt + answer**;
When `eos_early_stop = False`, generate() returns **answer**.
This caused issues during evaluation.
Proposed changes now ensure that the `generate()` function only returns **prompt + answer**, consistent with the output of generate() in llada and llada 1.5. (ps: In llada and llada 1.5, generate() returns **prompt + answer**)

Files changed (1) hide show
  1. modeling_llada2_moe.py +1 -1
modeling_llada2_moe.py CHANGED
@@ -1618,4 +1618,4 @@ class LLaDA2MoeModelLM(LLaDA2MoePreTrainedModel, GenerationMixin):
1618
  first_mask_position = mask_positions[0].item()
1619
  else:
1620
  first_mask_position = gen_length
1621
- return generated_answer[:, input_ids.shape[1] : input_ids.shape[1] + first_mask_position + 1]
 
1618
  first_mask_position = mask_positions[0].item()
1619
  else:
1620
  first_mask_position = gen_length
1621
+ return generated_answer[:, : input_ids.shape[1] + first_mask_position + 1]