Hi (@ayalaall ) Ayala,
I had a similar problem.   (I had a model in TF and I wanted to use BertViz).  Like you, I thought it should be possible to copy the TF weights into a Pytorch framework.  After all, they are just a bunch of numbers.   However, I couldn’t get it to work.   I could be wrong, but I don’t think there is an easy way to move a model from TF to Pytorch.
My solution (which worked eventually) was to start again and train a new model using Pytorch.   This made sense for me, as I wasn’t particularly expert in TF/keras, and I though it might be handy to learn Pytorch.
Afterwards, I thought it would have been Better to write a copy of BertViz that was designed to work with TF.   (How hard can it be…?).   If you are expert in TF and in Python then this might be a good solution for you.
A third possibility would be to look at the internal structures of the way the weights are stored for TF and for Pytorch, and to force your model’s numbers into a Pytorch-like structure.   Since both the TF and the Pytorch models are implementations of the same Attention-based Encoder, I think this should be theoretically possible.   It doesn’t sound easy though.
It is possible that somebody has written a TF-Pytorch Converter program.  I couldn’t find one when I looked, but that was nearly two years ago, so there might be one now.  If you ask another question on this forum with “TF toPytorch Model Conversion” in the title, somebody might know (but don’t hold your breath waiting).
It is possible that somebody has written a Visualisation tool for TF Bert models.  Again, I couldn’t find one two years ago.  How much have you searched?