What tool do you use to get the token prediction of each layer of large language models for Figure 2?
frankdarkluo opened this issue · comments
I simply use matplotlib to make a table for figure 2!
Thanks for the reply! But I am not asking about the drawing. I am curious how do you get the probability distribution from the middle (not output) layers? Thanks.
Just insert some code to the transformers package (modeling_llama.py and generation/utils.py) and get the predictions along the decoding steps. It makes the code ugly but it works. I didn't use any tools for that.