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]

The CI failing for Python 3.9 seems unrelated to this PR.

[qgallouedec]

The trl Trainers like PPO should not try to back propagate through the generated tokens, right?

Yes that's correct. The backprop is done on the output of a forward pass

[konrad-gerlach]

@qgallouedec Could you run the precommit to fix the linting issues? I haven't gotten it to work.

[konrad-gerlach]

I'm still working on adding some more tests and cleaning up the code a bit.