Add generation caching in TextEnvironment and fix bugs in TextEnvironment

[konrad-gerlach]

This PR mainly affects the TextEnvironment class and adds caching in between generation calls, in order to not have to recompute all previous activations when generating the next segment. This is mainly intended for use cases where many tool calls are performed sequentially and thus the activations for the (possibly quite large) system prompt would have to be calculated at each step. For stability, caching is optional.

Bug fixes:
This issue also addresses two bugs I encountered:

1. max_length checking in TextEnvironment class threw an error, as it assumed batching was present, when no batching existed.
   I fixed the bug and also added a check at generation time to ensure, that the padded inputs also do not exceed max length.
2. The StringStoppingCriteria did not take generated eos tokens into account, which I have now fixed.

RE testing:
I only made sure, that the tests in tests/test_environments.py were completing.
Using make test some tests were failing and the tests were taking a long time to run. However, the only tests, which call TextEnvironment seem to be in test_environments.py, so the rest should be unaffected as far as I know. Nevertheless, I would be grateful, if somebody else could run all the tests before merging. I suspect, that my environment may not be ideally configured. Is testing automated via a CI?

[konrad-gerlach]

I would be very grateful for a review by:
@lvwerra
@vwxyzjn
@younesbelkada
@qgallouedec
or any others, that feel up to the task.

[konrad-gerlach]

I was unable to execute the pre-commit hook, so I manually ran the linter.

[qgallouedec]

Thanks for the PR!
Let's see what's the CI outputs.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[konrad-gerlach]

Just to be sure, as I'm unfamiliar with their implementation: The trl Trainers like PPO should not try to back propagate through the generated tokens, right?

[konrad-gerlach]

I noticed, that GPT2 seems to only support the legacy cache format, so I am adding support for this,

[konrad-gerlach]

In the previous commits, I also fixed what I believe to be an off-by-one error in the StringStoppingCriteria: an inaccurate number for generated_tokens.

It appears, that it was still not fixed. Working on a solution and on testing StringStoppingCriteria

[konrad-gerlach]

As the tests did not include a encoder-decoder architecture, I did not test for it either. I think, that this is out of scope for this Pull Request. Where this was of concern in _generate_batched, I mirrored the implementation already provided.

[konrad-gerlach]

I added a TODO in the code, that the batch_size parameter should be exposed to the user. I think this is out of scope for this PR.

[konrad-gerlach]

Should there be a general warning for TextEnvironment, that support for Encoder Decoder is not (automatically) tested?

[konrad-gerlach]

@qgallouedec I am ready for review. Please also read the comments/ questions above and do the pre-commit hooks (as discussed). I hope you like the code :)

[konrad-gerlach]

A notable behavioural change of this PR is, that if input_ids exceeds TextEnvironment.max_length we truncate all sequences instead of allowing the model to error. This technically breaks the isolation of the samples within the batch, as a history might be truncated due to another history being too long, but I think this is acceptable behaviour for now. Future work might remove the limiting sequence(s), truncate their histories and continue generating using the remaining histories. This is out of scope for this PR.