stanfordnlp / pyvene

Stanford NLP Python Library for Understanding and Improving PyTorch Models via Interventions
http://pyvene.ai
Apache License 2.0
652 stars 63 forks source link

add support for dict outputs (select the first value) #178

Closed leogagnon closed 4 months ago

leogagnon commented 4 months ago

Description

Add simple (and naive) for dictionary output of modules. Similar to how tupless were dealt with, I just select the first value of the dict. Only made changes to BaseModel. I'm not yet super familiar with the code so there might be some other modification needed.

Testing Done

It works with my model which has module dict outputs (for ZeroIntervention and LoreftIntervention). Didn't do extensive testing as this felt quite straightforward.

Checklist:

aryamanarora commented 4 months ago

LGTM, pinging @frankaging: have you ever used modules which output dicts? Hasn't come up for me but this fix seems reasonable.

frankaging commented 4 months ago

LGTM! Later, we might want to further remove the assumption of intervening on the first element by creating field in the config to provide the intervening keyword.

frankaging commented 4 months ago

@leogagnon Thanks!