How do you set up an attention visualization tool in code to interpret and debug transformer model outputs