namin / llm-verified-with-monte-carlo-tree-search

LLM verified with Monte Carlo Tree Search
https://arxiv.org/abs/2402.08147
MIT License
252 stars 28 forks source link

Patch to handle EOS in limit_depth #57

Closed davidbrandfonbrener closed 7 months ago

davidbrandfonbrener commented 7 months ago

This patch handles EOS as if the model has hit the depth limit to prevent infinite expansions of the same string by attempting to generate after an EOS.

I changed the llm.py to (1) expose the eos_token and (2) set skip_special_tokens=False so that the EOS token will be added into the strings returned by gen().

Then in common.py, I check for eos_token at the same time as checking for max_depth, and prevent further expansion of this node is EOS has been reached.