Refactor hooks / Save Attention Maps script #34

v0xie · 2024-05-13T03:02:27Z

Adds module_hooks.py with functions to add fields / forward hooks to modules
Adds prompt_utils with a function to get token count.
Adds Save Attention Maps which uses the new functions in module_hooks to do the hooking
Save Attention Maps and other debug scripts will be loaded iff the INCANT_DEBUG environment variable is set.

v0xie added 14 commits May 2, 2024 15:27

add module_hooks.py

a9656f5

add function to add forward hook

df35a5a

remove self arg from modules_add_field

c0007e8

add save_attn_maps

1766d21

update save_attn_maps script

1f0c1f3

update saving logic

063d848

remove warning if removing non-existent field

c73d5cb

saving attention maps

ae34cb4

multiply by text embeddings for cross-attention

b48a995

add prompt_utils.py

47fc67a

update save format

93cfe63

change default module_name_filter value

a8cbe9a

fix saveattnmaps postprocess running without token count

f0474a7

load saveattnmaps script when INCANT_DEBUG env set

603c9cd

v0xie merged commit 6aa847f into dev May 13, 2024

v0xie mentioned this pull request May 14, 2024

S-CFG, Optimizations, and More #37

Merged

Provide feedback