How to plot bert model? (Transfomer models)

[indramal]

I try to plot bert model using this package. But I unable to do it.

Code:

from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hello world!", return_tensors="pt")

After that how to plot it?

import hiddenlayer as hl
hl.build_graph(model, inputs[0])

[Foxglove144]

You need to extract input_ids and attention_mask from the tokenizers.Encoding object and pass them as a tuple to the hl.build_graph function.

import hiddenlayer as hl
from transformers import AutoModel, AutoTokenizer

model = AutoModel.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hello world!", return_tensors="pt")

input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]

hl.build_graph(model, (input_ids, attention_mask))