TextEnvironments

[lvwerra]

This PR adds multi-turn text environment to TRL.

Target API

env = Environment(model, tokenizer, prompt, tools)

for tasks in ppo_trainer.dataloader:
    histories = env.run(tasks)
    tokens, mask = histories.get_tokens()
    ppo_trainer.step(tokens, mask=mask, histories.rewards)
    
# alternatively, this would probably be much nicer
for tasks in ppo_trainer.dataloader:
    tokens, masks, rewards, history = env.run(tasks)
    ppo_trainer.step(tokens, masks, rewards)

Todos

• text to token+mask conversion | add v1 masking #429
• penalty for tool calls
• masking in PPOTrainer after fwd pass | add v1 masking #429
• end-to-end example calculater

Current working example

from trl import TextEnvironment, TextHistory, AutoModelForCausalLMWithValueHead
from transformers import AutoModelForCausalLM, AutoTokenizer, load_tool

tool = load_tool("ybelkada/simple-calculator")

model_id = "gpt2-xl"

model = AutoModelForCausalLMWithValueHead.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.pad_token = tokenizer.eos_token

prompt = """\
What is 12.1 + 1 - 3?
<request><SimpleCalculatorTool>12.1 + 1<call>13.1<response>
<request><SimpleCalculatorTool>13.1 - 3<call>10.1<response>
Result = 10.1 <submit>

"""

reward_fn = lambda x: 1

env = TextEnvironment(model, tokenizer,[tool], reward_fn, prompt, generation_kwargs={"max_new_tokens": 32})
h = env.run(["What is 387 * 228?"])

h[0].show()

Result:

What is 12.1 + 1 - 3?
<request><SimpleCalculatorTool>12.1 + 1<call>13.1<response>
<request><SimpleCalculatorTool>13.1 - 3<call>10.1<response>
Result = 10.1 <submit>

What is 387 * 228?

<request><SimpleCalculatorTool>387 * 228<call>88236.0<response>

Result = 88236.0 <submit>
Reward: 1

Where the 88236.0<response> segement was generated by the tool call.

[HuggingFaceDocBuilderDev]

The documentation is not available anymore as the PR was closed or merged.

[vwxyzjn]

Thanks @younesbelkada for the comments, I have addressed most of them. @lvwerra the PR looks good. Quick question: do we want to merge docs/source/learning_tools.mdx and docs/source/text_environments.md?

[lvwerra]

Quick question: do we want to merge docs/source/learning_tools.mdx and docs/source/text_environments.md?

I thought we could have text_environments.md as the basic doc for how TextEnvs work and the learning_tools.mdx as the more hands-on guide also linking the experiment scripts. Wdyt?

Also @vwxyzjn would you mind fixing the scripts so they pass the quality checks. In my opinion we can also exclude them from the quality tests, no strong opinion.

[younesbelkada]

Looking great ! Thanks for all your efforts, left 4 nits, otherwise LGTM !