feat: streaming response for HookedTransformer.generate

[hijohnnylin]

Description

This adds a stream_output flag (default False) to HookedTransformer.generate which streams the tokens as they're generated, instead of waiting for the whole result to be generated first.

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[hijohnnylin]

still a draft as I have not written docs, tests, etc.

[hijohnnylin]

the current implementation breaks existing calls where stream_output is false because yield forces the method to become a Generator even if that code path is not used. for now just copying the method.

[bryce13950]

Very nice! Looking forward to messing around with this. I will keep an eye out for a test, and then we can get it into a build right away.