[pip package] feature request: pipeline.generate: add ability to get the state, if it was not provided
Author: MaykeyeCreated Mar 18, 2024Updated Mar 26, 2024
Right now if I want to have state after the generation, PIPELINE.generate doesn't allow it if I didn't have it already: default value is None
So I need to do something silly like generate a single token with model.forward and pass state from there or make model wrapper which has forward and remembers the state.
So possible solutions to get the model state from generate:
- Add
return_stateargument and be like
if return_state: return out_str, state
return out_str Major drawback: it breaks clear simple API as now there are two return types, which will make LSPs/linters/etc unhappy and they will complain
- Add argument
callback_with_state: booland if it's true callcallback(tmp, state)rather thancallback(tmp)
Drawback: it's not always called, but I'm not sure it's that big deal
- Rename function to
generate_with_statewhich returns out_str and state.generatereplace with agenerate_with_state(...)[0]
This will keep the existing API
- Store the state. Make a field
PIPELINE.last_stateand storestatethere.
Wastes memory when is not always desirable.
Source: BlinkDL/ChatRWKV