Code execution & tool use

[KiddoZhu]

Add a basic version of code execution & tool use (#55, #56)

Implemented by a custom LogitProcessor in vllm backend. I implement tools as pre-defined variables in code environment so that there is no need to distinguish code and tool use. It also makes maintenance easier. The only concern is that code execution must be enabled whenever we want to enable tool use, though we can avoid telling the existence of arbitrary code execution beyond the tools in the prompt.
 

Supported features

• On-the-fly multiple tool calls.
• Compatible with batch decoding.
• Stateful code executor. In the same generation, functions and variables will be passed to the next code snippet.

Examples (test_vllm_tools.py)

<code>x = 3; y = 4</code>
This is some regular text.
<code>x + y</code>
<result>7</result>

<code>retrieve('Jen-Hsun Huang')</code>

<result>
['Nvidia was established in 1993 by Jen-Hsun Huang, Curtis Priem, and Chris '
 'Malachowsky. In 2000 Nvidia took intellectual possession of 3dfx, one of the '
 'biggest GPU producers in 1990s.']
</result>

Tokenizer issue

Although I design CodeLogitProcessor to be tokenizer-agnostic as much as possible, there may be side cases where tokens don't split exactly at the end of </code>. For example, the tokenizer of GPT-4o will generate the following

... </##code##>x

If we tweak >x to be >, it will change the log prob in RL. My current solution is to not touch any generated token and directly append results afterwards

... </code>x<result> ...

[KiddoZhu]

Now we have an implementation based on stop strings, making it independent of vllm and hf backends.

The function generate_with_code_and_tools(policy, input_batch, tokenizer) can be used as a drop-in replacement of policy.generate(input_batch). It supports

• Multiple code execution & tool use. It will continue generation until hitting an EOS token or user-specified stop strings.
• Batch code execution based on ray.remote (@SahilJain314 Could you please check if we need to take special care of the ray workers here?)

@SahilJain314 The way we calculated generated_lengths in HFPolicyWorker is buggy when the generation is stopped by custom stop strings (e.g. </code> in code execution). It will count one more EOS token after the stop strings, since EOS token is also used for padding. I changed it to count the nonzero entries of generation_logprobs, assuming that any generated logprob is different from the padding value 0. Please let me know if you foresee any issue with this implementation.

I'll merge the main branch and clean up the logit processor implementation.

[KiddoZhu]

Done. Ready for review.

[SahilJain314]

@terrykong afaik this shouldn't be gitignored, right?

[terrykong]

yea, we should not ignore the lock

[SahilJain314]

This assumption doesn't hold. Generated logprobs are actually surprisingly sometimes 0.

[KiddoZhu]

Then how can we safely decide the length? The old implementation mistakes padded eos token as generated eos token. It's a serious problem for custom stop strings.

[SahilJain314]

Can we replace this with the existing run_multi_turn_generation function? and use environments to handle the stop tokens as we do right now?

[KiddoZhu]

Sounds good to me. Will code execution always be invoked in a multi-turn chat format? How about changing run_multi_turn_rollout to normal prompt format, and then provide an additional interface on top of that to support chat messages?

[SahilJain314]

the default multi turn chat actually runs without chat templating (the template is applied in the string/tokens directly). To chat template it, the environment would have to do it. (Open to suggestions on this approach)

[SahilJain314]

I'm a little hesistant to push this out without filesystem isolation

[KiddoZhu]

I'm a little hesistant to push this out without filesystem isolation

Good point! What's the solution in your mind? I can think of chroot or docker, but they are too heavy. If we believe the code is just untrusted, not malicious, a simple solution is to 1) run the code in a temporary directory 2) override open() to deny access beyond the temporary directory 3) override builtins.__import__ to deny modules like os, sys and subprocess.

[KiddoZhu]

Sorry I messed up the commit history. This branch will be no longer touched but kept for retrieving history. I will open a new branch and pull request.