Fix hook counters for non-generation calls.

stanfordnlp / pyvene

Stanford NLP Python Library for Understanding and Improving PyTorch Models via Interventions

http://pyvene.ai

Apache License 2.0

609 stars 59 forks source link

Fix hook counters for non-generation calls. #16

Closed jhuang81 closed 10 months ago

jhuang81 commented 10 months ago

Fix hook counters for non-generation calls by adding a flag to indicate whether we are in the "generate" call or not.

Similar to the hooks, this implementation can have unintended behavior if the program crashes in the middle of the generate call, which will leave the _is_generation to True. It would be great if we could find better indicators for the "generate" call.