Thanks for your response… I am just giving example.
import tensorflow as tf
from transformers import BertTokenizerFast
from transformers import TFBertModel
mlm_model = TFBertModel.from_pretrained(‘bert-base-uncased’,output_attentions=True)
input_word_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“input_word_ids”)
token_type_ids = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“token_type_ids”)
attention_mask = tf.keras.layers.Input(shape=(50,), dtype=tf.int32,name=“attention_mask”)
mlm_out=mlm_model([input_word_ids,token_type_ids,attention_mask])
print(“Attention Heads are”,mlm_out[2])
This mlm_out[2] has 12 attention heads now I want to extract the one attention head from mlm_out[2] which will have the highest probability. Is it possible or not …