<commit_before># coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Change log
# 12/12/18  <NAME>   Adapted to BERT model
# 12/19/18  <NAME>   Assorted cleanup. Changed orientation of attention matrices. Updated comments.
# 12/30/20  <NAME>   Enable multiple visualizations in one notebook
# 01/16/21  <NAME>   Dark mode


"""Module for postprocessing and displaying transformer attentions.

This module is designed to be called from an ipython notebook.
"""

import json
import os
import uuid
from collections import defaultdict

import torch
from IPython.core.display import display, HTML, Javascript


def show(model, model_type, tokenizer, sentence_a, sentence_b=None, display_mode='dark', layer=None, head=None, html_action='view'):

    # Generate unique div id to enable multiple visualizations in one notebook

    if sentence_b:
        vis_html = """
        <div id="bertviz" style='padding:8px'>
          <span style="user-select:none">
            <span class="dropdown-label">Layer: </span><select id="layer"></select>
            <span class="dropdown-label">Head: </span><select id="att_head"></select>
            <span class="dropdown-label">Attention: </span><select id="filter">
              <option value="all">All</option>
              <option value="aa">Sentence A -> Sentence A</option>
              <option value="ab">Sentence A -> Sentence B</option>
              <option value="ba">Sentence B -> Sentence A</option>
              <option value="bb">Sentence B -> Sentence B</option>
            </select>
          </span>
          <div id='vis'></div>
        </div>
        """
    else:
        vis_html = """
            <div id="bertviz" style='padding:8px'>
              <span style="user-select:none">
                <span class="dropdown-label">Layer: </span><select id="layer"></select>
                <span class="dropdown-label">Head: </span> <select id="att_head"></select>
              </span>
              <div id='vis'></div>
            </div>
         """

    # require.js must be imported for Colab or JupyterLab:
    if html_action == 'view'
        display(HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>'))
        display(HTML(vis_html))
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
        attn_data = get_attention(model, model_type, tokenizer, sentence_a, sentence_b, include_queries_and_keys=True)
        if model_type == 'gpt2':
            bidirectional = False
        else:
            bidirectional = True
        params = {
            'attention': attn_data,
            'default_filter': "all",
            'bidirectional': bidirectional,
            'display_mode': display_mode,
            'layer': layer,
            'head': head
        }
        vis_js = open(os.path.join(__location__, 'neuron_view.js')).read()
        display(Javascript('window.bertviz_params = %s' % json.dumps(params)))
        display(Javascript(vis_js))

    elif html_action == 'return':
        html1 = HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>')
        html2 = HTML(vis_html)
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))<commit_msg>add support for 'html_action' at neuron_view<commit_after># coding=utf-8
# Copyright 2018 The Tensor2Tensor Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Change log
# 12/12/18  <NAME>   Adapted to BERT model
# 12/19/18  <NAME>   Assorted cleanup. Changed orientation of attention matrices. Updated comments.
# 12/30/20  <NAME>   Enable multiple visualizations in one notebook
# 01/16/21  <NAME>   Dark mode


"""Module for postprocessing and displaying transformer attentions.

This module is designed to be called from an ipython notebook.
"""

import json
import os
import uuid
from collections import defaultdict

import torch
from IPython.core.display import display, HTML, Javascript


def show(model, model_type, tokenizer, sentence_a, sentence_b=None, display_mode='dark', layer=None, head=None, html_action='view'):

    # Generate unique div id to enable multiple visualizations in one notebook

    if sentence_b:
        vis_html = """
        <div id="bertviz" style='padding:8px'>
          <span style="user-select:none">
            <span class="dropdown-label">Layer: </span><select id="layer"></select>
            <span class="dropdown-label">Head: </span><select id="att_head"></select>
            <span class="dropdown-label">Attention: </span><select id="filter">
              <option value="all">All</option>
              <option value="aa">Sentence A -> Sentence A</option>
              <option value="ab">Sentence A -> Sentence B</option>
              <option value="ba">Sentence B -> Sentence A</option>
              <option value="bb">Sentence B -> Sentence B</option>
            </select>
          </span>
          <div id='vis'></div>
        </div>
        """
    else:
        vis_html = """
            <div id="bertviz" style='padding:8px'>
              <span style="user-select:none">
                <span class="dropdown-label">Layer: </span><select id="layer"></select>
                <span class="dropdown-label">Head: </span> <select id="att_head"></select>
              </span>
              <div id='vis'></div>
            </div>
         """

    # require.js must be imported for Colab or JupyterLab:
    if html_action == 'view':
        display(HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>'))
        display(HTML(vis_html))
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
        attn_data = get_attention(model, model_type, tokenizer, sentence_a, sentence_b, include_queries_and_keys=True)
        if model_type == 'gpt2':
            bidirectional = False
        else:
            bidirectional = True
        params = {
            'attention': attn_data,
            'default_filter': "all",
            'bidirectional': bidirectional,
            'display_mode': display_mode,
            'layer': layer,
            'head': head
        }
        vis_js = open(os.path.join(__location__, 'neuron_view.js')).read()
        display(Javascript('window.bertviz_params = %s' % json.dumps(params)))
        display(Javascript(vis_js))

    elif html_action == 'return':
        html1 = HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>')
        html2 = HTML(vis_html)
        __location__ = os.path.realpath(
            os.path.join(os.getcwd(), os.path.dirname(__file__)))
{'fraction_non_alphanumeric': 0.12916880451513596, 'fraction_numerical': 0.011544381734222679, 'mean_word_length': 2.730622009569378, 'pattern_counts': {'":': 0, '<': 103, '<?xml version=': 0, '>': 111, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2188088', 'n_tokens_mistral': 2331, 'n_tokens_neox': 2231, 'n_words': 614}