BoN sampling: ArrowInvalid: offset overflow while concatenating arrays

[cassidylaidlaw]

After training a reward model according to the instructions, I tried to run BoN sampling via the following command:

python src/bon/run_bon_pipeline.py models/rm-pythia-44m_seed0 --seeds 0

However, I got the following traceback: bon.log.

I was able to fix this error using a workaround in this issue. In run_bon_pipeline.py, I changed

gold_labelled_generations = gold_labelled_generations.map(_truncate_answers)

to

gold_labelled_generations = gold_labelled_generations.map(
    _truncate_answers,
    writer_batch_size=8,
)

It might be good to update the code based on this since others could have the same issue.

[tsWen0309]

Thanks! That works for me too.