This is the open-source code for paper:
Retrieval Head Mechanistically Explains Long-Context Factuality.
This code is implemented based on Needle In a HayStack.
【Update】 Support Phi3 now, thanks to the contribution made by @Wangmerlyn.
An algorithm that statistically calculate the retrieval score of attention heads in a transformer model.
Because FlashAttention can not return attention matrix, this algorithm is implemented by… See the full description on the dataset page:
https://huggingface.co/datasets/shaswat123/NIAH.
