Hi @Rasputin312 !

I was able to render the visualization with bertviz library.

The default moview_view html_action is view that does not work with Databricks notebook. Instead, using the returned HTML, we can visualize the model.

display(model_view(attention, tokens, html_action="return")) 

From model_view doc:

html_action: Specifies the action to be performed with the generated HTML object - 'view' (default): Displays the generated HTML representation as a notebook cell output - 'return' : Returns an HTML object containing the generated view for further processing or custom visualization

The entire notebook code is available here. https://gist.github.com/koji-kawamura-db/49c52d5d96b25b936e038fc21d3bbcec