Merge pull request jessevig#1 from cyber-raskolnikov/feature/neuron_v…

…iew_html_return Feature/neuron view html return
yuanzhongqiao · Mar 23, 2022 · 2ec7128 · 2ec7128
2 parents aefce0c + 10508a1
commit 2ec7128
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 24 deletions.
diff --git a/bertviz/model_view.py b/bertviz/model_view.py
@@ -247,5 +247,4 @@ def model_view(
         return head_html
 
     else:
-        raise ValueError("'html_action' parameter must be 'view' or 'return")
-
+        raise ValueError("'html_action' parameter must be 'view' or 'return")
diff --git a/bertviz/neuron_view.js b/bertviz/neuron_view.js
@@ -10,12 +10,13 @@
  * 01/01/21  Jesse Vig   Change to bertviz-specific naming conventions so as not to interfere with other html elements
  * 01/16/21  Jesse Vig   Dark mode
  * 02/06/21  Jesse Vig   Move require config from separate jupyter notebook step
+ * 03/23/22  Daniel SC   Update requirement URLs for d3 and jQuery (source of bug not allowing end result to be displayed on browsers)
  **/
 
 require.config({
   paths: {
-      d3: '//cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min',
-    jquery: '//ajax.googleapis.com/ajax/libs/jquery/2.0.0/jquery.min',
+      d3: 'https://cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min',
+    jquery: 'https://cdnjs.cloudflare.com/ajax/libs/jquery/2.0.0/jquery.min',
   }
 });
 

diff --git a/bertviz/neuron_view.py b/bertviz/neuron_view.py
@@ -34,7 +34,7 @@
 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):
+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
 
@@ -65,27 +65,58 @@ def show(model, model_type, tokenizer, sentence_a, sentence_b=None, display_mode
               <div id='vis'></div>
             </div>
          """
+
     # require.js must be imported for Colab or JupyterLab:
-    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
+    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__)))
+        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()
+
+        script1 = '\n<script type="text/javascript">\n' + Javascript('window.bertviz_params = %s' % json.dumps(params)).data + '\n</script>\n'
+        script2= '\n<script type="text/javascript">\n' + Javascript(vis_js).data + '\n</script>\n'
+
+        neuron_html = HTML(html1.data + html2.data + script1 + script2)
+        return neuron_html
+
     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))
+        raise ValueError("'html_action' parameter must be 'view' or 'return")
 
 
 def get_attention(model, model_type, tokenizer, sentence_a, sentence_b=None, include_queries_and_keys=False):