How I can extract attention map?

facebookresearch / esm

Evolutionary Scale Modeling (esm): Pretrained language models for proteins

MIT License

3.16k stars 627 forks source link

How I can extract attention map? #55

Closed nasserhashemi closed 3 years ago

nasserhashemi commented 3 years ago

Hi there; I hope you are well; I have question, given a sequence as input in esm1b, how I can extract 660 attention map associated with each head in each layer? Thanks so much

tomsercu commented 3 years ago

Hi Nasser, model.forward() has an argument need_head_weights, link. The return dictionary will then contain result["attentions"] see L183. Feel free to re-open if you have any issues!