Fix bug: Prevent inconsistent return value with `eos_early_stop`
Browse filesWhen `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**)
- 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[:,
|
|
|
|
| 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]
|