Spaces:

gabrielanicole
/

makenmtviz

Sleeping

App Files Files Community

Gabriela Nicole Gonzalez Saez commited on Feb 15

Commit

056bbdc

•

1 Parent(s): fc37a00

Add files

Browse files

Files changed (4) hide show

app.py +114 -0
bertviz_gradio.py +248 -0
plotsjs_bertviz.js +430 -0
requirements.txt +2 -0

app.py ADDED Viewed

	@@ -0,0 +1,114 @@

+import gradio as gr
+import inseq
+import captum
+import torch
+import os
+# import nltk
+import argparse
+import random
+import numpy as np
+from argparse import Namespace
+from tqdm.notebook import tqdm
+from torch.utils.data import DataLoader
+from functools import partial
+from transformers import AutoTokenizer, MarianTokenizer, AutoModel, AutoModelForSeq2SeqLM, MarianMTModel
+from bertviz import model_view, head_view
+from bertviz_gradio import head_view_mod
+def get_bertvis_data(input_text, lg_model):
+	tokenizer_tr = dict_tokenizer_tr[lg_model]
+	model_tr = dict_models_tr[lg_model]
+	input_ids = tokenizer_tr(input_text, return_tensors="pt", padding=True)
+	result_att = model_tr.generate(**input_ids,
+		return_dict_in_generate=True,
+		output_attentions =True,
+		output_scores=True,
+	)
+	# tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0])
+	# tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0])
+	tgt_text = tokenizer_tr.decode(result_att.sequences[0], skip_special_tokens=True)
+	print(tgt_text)
+	outputs = model_tr(input_ids=input_ids.input_ids,
+					decoder_input_ids=result_att.sequences,
+					output_attentions =True,
+					)
+	print(tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]))
+	# print(tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0]), tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]))
+	html_attentions = head_view_mod(
+		encoder_attention = outputs.encoder_attentions,
+		cross_attention = outputs.cross_attentions,
+		decoder_attention = outputs.decoder_attentions,
+		encoder_tokens = tokenizer_tr.convert_ids_to_tokens(input_ids.input_ids[0]),
+		decoder_tokens = tokenizer_tr.convert_ids_to_tokens(result_att.sequences[0]),
+		html_action='gradio'
+		)
+	return html_attentions, tgt_text
+## First create html and divs
+html = """
+<html>
+<script async src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>
+<script src="https://cdnjs.cloudflare.com/ajax/libs/jquery/2.0.0/jquery.min"></script>
+<script async data-require="d3@3.5.3" data-semver="3.5.3" src="//cdnjs.cloudflare.com/ajax/libs/d3/3.5.3/d3.js"></script>
+  <body>
+    <div id="bertviz"></div>
+    <div id="d3_beam_search"></div>
+  </body>
+</html>
+"""
+def sentence_maker(w1, model, var2={}):
+  #translate and get internal values
+  params,tgt = get_bertvis_data(w1, model)
+  ### get translation
+  return [tgt, params['params'],params['html2'].data]
+def sentence_maker2(w1,j2):
+  #  json_value = {'one':1}
+  #  return f"{w1['two']} in sentence22..."
+   print(w1,j2)
+   return "in sentence22..."
+with gr.Blocks(js="plotsjs_bertviz.js") as demo:
+    gr.Markdown("""
+              # MAKE NMT Workshop \t `BertViz` \n
+              https://github.com/jessevig/bertviz
+              """)
+    with gr.Row():
+      with gr.Column(scale=1):
+        in_text = gr.Textbox(label="Source Text")
+        out_text  = gr.Textbox(label="Target Text")
+        out_text2  = gr.Textbox(visible=False)
+        var2 = gr.JSON(visible=False)
+        btn = gr.Button("Create sentence.")
+        radio_c = gr.Radio(choices=['en-zh', 'en-es', 'en-fr'], value="en-zh", label= '', container=False)
+      with gr.Column(scale=4):
+        gr.Markdown("Attentions: ")
+        input_mic = gr.HTML(html)
+        out_html = gr.HTML()
+      btn.click(sentence_maker, [in_text,radio_c], [out_text,var2,out_html], js="(in_text,radio_c) => testFn_out(in_text,radio_c)") #should return an output comp.
+      out_text.change(sentence_maker2, [out_text, var2], out_text2, js="(out_text,var2) => testFn_out_json(var2)") #
+      # out_text.change(sentence_maker2, [out_text, var2], out_text2, js="(out_text,var2) => testFn_out_json(var2)") #
+    # run script function on load,
+    # demo.load(None,None,None,js="plotsjs.js")
+if __name__ == "__main__":
+    demo.launch()

bertviz_gradio.py ADDED Viewed

	@@ -0,0 +1,248 @@

+import json
+import os
+import uuid
+from IPython.core.display import display, HTML, Javascript
+from bertviz.util import format_special_chars, format_attention, num_layers
+def head_view_mod(
+        attention=None,
+        tokens=None,
+        sentence_b_start=None,
+        prettify_tokens=True,
+        layer=None,
+        heads=None,
+        encoder_attention=None,
+        decoder_attention=None,
+        cross_attention=None,
+        encoder_tokens=None,
+        decoder_tokens=None,
+        include_layers=None,
+        html_action='view'
+):
+    """Render head view
+        Args:
+            For self-attention models:
+                attention: list of ``torch.FloatTensor``(one for each layer) of shape
+                    ``(batch_size(must be 1), num_heads, sequence_length, sequence_length)``
+                tokens: list of tokens
+                sentence_b_start: index of first wordpiece in sentence B if input text is sentence pair (optional)
+            For encoder-decoder models:
+                encoder_attention: list of ``torch.FloatTensor``(one for each layer) of shape
+                    ``(batch_size(must be 1), num_heads, encoder_sequence_length, encoder_sequence_length)``
+                decoder_attention: list of ``torch.FloatTensor``(one for each layer) of shape
+                    ``(batch_size(must be 1), num_heads, decoder_sequence_length, decoder_sequence_length)``
+                cross_attention: list of ``torch.FloatTensor``(one for each layer) of shape
+                    ``(batch_size(must be 1), num_heads, decoder_sequence_length, encoder_sequence_length)``
+                encoder_tokens: list of tokens for encoder input
+                decoder_tokens: list of tokens for decoder input
+            For all models:
+                prettify_tokens: indicates whether to remove special characters in wordpieces, e.g. Ġ
+                layer: index (zero-based) of initial selected layer in visualization. Defaults to layer 0.
+                heads: Indices (zero-based) of initial selected heads in visualization. Defaults to all heads.
+                include_layers: Indices (zero-based) of layers to include in visualization. Defaults to all layers.
+                    Note: filtering layers may improve responsiveness of the visualization for long inputs.
+                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
+    """
+    attn_data = []
+    if attention is not None:
+        if tokens is None:
+            raise ValueError("'tokens' is required")
+        if encoder_attention is not None or decoder_attention is not None or cross_attention is not None \
+                or encoder_tokens is not None or decoder_tokens is not None:
+            raise ValueError("If you specify 'attention' you may not specify any encoder-decoder arguments. This"
+                             " argument is only for self-attention models.")
+        if include_layers is None:
+            include_layers = list(range(num_layers(attention)))
+        attention = format_attention(attention, include_layers)
+        if sentence_b_start is None:
+            attn_data.append(
+                {
+                    'name': None,
+                    'attn': attention.tolist(),
+                    'left_text': tokens,
+                    'right_text': tokens
+                }
+            )
+        else:
+            slice_a = slice(0, sentence_b_start)  # Positions corresponding to sentence A in input
+            slice_b = slice(sentence_b_start, len(tokens))  # Position corresponding to sentence B in input
+            attn_data.append(
+                {
+                    'name': 'All',
+                    'attn': attention.tolist(),
+                    'left_text': tokens,
+                    'right_text': tokens
+                }
+            )
+            attn_data.append(
+                {
+                    'name': 'Sentence A -> Sentence A',
+                    'attn': attention[:, :, slice_a, slice_a].tolist(),
+                    'left_text': tokens[slice_a],
+                    'right_text': tokens[slice_a]
+                }
+            )
+            attn_data.append(
+                {
+                    'name': 'Sentence B -> Sentence B',
+                    'attn': attention[:, :, slice_b, slice_b].tolist(),
+                    'left_text': tokens[slice_b],
+                    'right_text': tokens[slice_b]
+                }
+            )
+            attn_data.append(
+                {
+                    'name': 'Sentence A -> Sentence B',
+                    'attn': attention[:, :, slice_a, slice_b].tolist(),
+                    'left_text': tokens[slice_a],
+                    'right_text': tokens[slice_b]
+                }
+            )
+            attn_data.append(
+                {
+                    'name': 'Sentence B -> Sentence A',
+                    'attn': attention[:, :, slice_b, slice_a].tolist(),
+                    'left_text': tokens[slice_b],
+                    'right_text': tokens[slice_a]
+                }
+            )
+    elif encoder_attention is not None or decoder_attention is not None or cross_attention is not None:
+        if encoder_attention is not None:
+            if encoder_tokens is None:
+                raise ValueError("'encoder_tokens' required if 'encoder_attention' is not None")
+            if include_layers is None:
+                include_layers = list(range(num_layers(encoder_attention)))
+            encoder_attention = format_attention(encoder_attention, include_layers)
+            attn_data.append(
+                {
+                    'name': 'Encoder',
+                    'attn': encoder_attention.tolist(),
+                    'left_text': encoder_tokens,
+                    'right_text': encoder_tokens
+                }
+            )
+        if decoder_attention is not None:
+            if decoder_tokens is None:
+                raise ValueError("'decoder_tokens' required if 'decoder_attention' is not None")
+            if include_layers is None:
+                include_layers = list(range(num_layers(decoder_attention)))
+            decoder_attention = format_attention(decoder_attention, include_layers)
+            attn_data.append(
+                {
+                    'name': 'Decoder',
+                    'attn': decoder_attention.tolist(),
+                    'left_text': decoder_tokens,
+                    'right_text': decoder_tokens
+                }
+            )
+        if cross_attention is not None:
+            if encoder_tokens is None:
+                raise ValueError("'encoder_tokens' required if 'cross_attention' is not None")
+            if decoder_tokens is None:
+                raise ValueError("'decoder_tokens' required if 'cross_attention' is not None")
+            if include_layers is None:
+                include_layers = list(range(num_layers(cross_attention)))
+            cross_attention = format_attention(cross_attention, include_layers)
+            attn_data.append(
+                {
+                    'name': 'Cross',
+                    'attn': cross_attention.tolist(),
+                    'left_text': decoder_tokens,
+                    'right_text': encoder_tokens
+                }
+            )
+    else:
+        raise ValueError("You must specify at least one attention argument.")
+    if layer is not None and layer not in include_layers:
+        raise ValueError(f"Layer {layer} is not in include_layers: {include_layers}")
+    # Generate unique div id to enable multiple visualizations in one notebook
+    # vis_id = 'bertviz-%s'%(uuid.uuid4().hex)
+    vis_id = 'bertviz'#-%s'%(uuid.uuid4().hex)
+    # Compose html
+    if len(attn_data) > 1:
+        options = '\n'.join(
+            f'<option value="{i}">{attn_data[i]["name"]}</option>'
+            for i, d in enumerate(attn_data)
+        )
+        select_html = f'Attention: <select id="filter">{options}</select>'
+    else:
+        select_html = ""
+    vis_html = f"""
+        <div id="{vis_id}" style="font-family:'Helvetica Neue', Helvetica, Arial, sans-serif;">
+            <span style="user-select:none">
+                Layer: <select id="layer"></select>
+                {select_html}
+            </span>
+            <div id='vis'></div>
+        </div>
+    """
+    for d in attn_data:
+        attn_seq_len_left = len(d['attn'][0][0])
+        if attn_seq_len_left != len(d['left_text']):
+            raise ValueError(
+                f"Attention has {attn_seq_len_left} positions, while number of tokens is {len(d['left_text'])} "
+                f"for tokens: {' '.join(d['left_text'])}"
+            )
+        attn_seq_len_right = len(d['attn'][0][0][0])
+        if attn_seq_len_right != len(d['right_text']):
+            raise ValueError(
+                f"Attention has {attn_seq_len_right} positions, while number of tokens is {len(d['right_text'])} "
+                f"for tokens: {' '.join(d['right_text'])}"
+            )
+        if prettify_tokens:
+            d['left_text'] = format_special_chars(d['left_text'])
+            d['right_text'] = format_special_chars(d['right_text'])
+    params = {
+        'attention': attn_data,
+        'default_filter': "0",
+        'root_div_id': vis_id,
+        'layer': layer,
+        'heads': heads,
+        'include_layers': include_layers
+    }
+    # require.js must be imported for Colab or JupyterLab:
+    if html_action == 'gradio':
+        html1 = HTML('<script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js"></script>')
+        html2 = HTML(vis_html)
+        return {'html1': html1, 'html2' : html2, 'params': params }
+    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__)))
+        vis_js = open(os.path.join(__location__, 'head_view.js')).read().replace("PYTHON_PARAMS", 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__)))
+        vis_js = open(os.path.join(__location__, 'head_view.js')).read().replace("PYTHON_PARAMS", json.dumps(params))
+        html3 = Javascript(vis_js)
+        script = '\n<script type="text/javascript">\n' + html3.data + '\n</script>\n'
+        head_html = HTML(html1.data + html2.data + script)
+        return head_html
+    else:
+        raise ValueError("'html_action' parameter must be 'view' or 'return")

plotsjs_bertviz.js ADDED Viewed

	@@ -0,0 +1,430 @@

+async () => {
+	// set testFn() function on globalThis, so you html onlclick can access it
+	 globalThis.testFn = () => {
+	   document.getElementById('demo').innerHTML = "Hello-bertviz?"
+	 };
+	//  await import * as mod from "/my-module.js";
+	 const d3 = await import("https://cdn.jsdelivr.net/npm/d3@5/+esm");
+	 const $ = await import("https://cdn.jsdelivr.net/npm/jquery@3.7.1/dist/jquery.min.js");
+	 globalThis.$ = $;
+	//  const $ = await import("https://cdn.jsdelivr.net/npm/jquery@2/+esm");
+	// import $ from "jquery";
+	// import * as d3 from "https://cdn.jsdelivr.net/npm/d3@7/+esm";
+	//  await import("https://cdn.jsdelivr.net/npm/jquery@2/+esm");
+	// export for others scripts to use
+	// window.$ = window.jQuery = jQuery;
+	// const d3 = await import("https://cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min");
+	// const $ = await import("https://cdnjs.cloudflare.com/ajax/libs/jquery/2.0.0/jquery.min");
+	globalThis.d3Fn = () => {
+		d3.select('#viz').append('svg')
+				.append('rect')
+				.attr('width', 50)
+				.attr('height', 50)
+				.attr('fill', 'black')
+				.on('mouseover', function(){d3.select(this).attr('fill', 'red')})
+				.on('mouseout', function(){d3.select(this).attr('fill', 'black')});
+    };
+	//
+	globalThis.testFn_out = (val,model) => {
+		// document.getElementById('demo').innerHTML = val
+		console.log(val);
+		// globalThis.d3Fn();
+		return([val,model]);
+	  };
+	globalThis.testFn_out_json = (data) => {
+		console.log(data);
+		var $ = jQuery;
+		console.log($('#viz'));
+		attViz(data);
+		return(['string', {}])
+	};
+	function attViz(PYTHON_PARAMS) {
+		var $ = jQuery;
+	const params = PYTHON_PARAMS; // HACK: PYTHON_PARAMS is a template marker that is replaced by actual params.
+	const TEXT_SIZE = 15;
+	const BOXWIDTH = 110;
+	const BOXHEIGHT = 22.5;
+	const MATRIX_WIDTH = 115;
+	const CHECKBOX_SIZE = 20;
+	const TEXT_TOP = 30;
+	console.log("d3 version in ffuntions", d3.version)
+	let headColors;
+	try {
+		headColors = d3.scaleOrdinal(d3.schemeCategory10);
+	} catch (err) {
+		console.log('Older d3 version')
+		headColors = d3.scale.category10();
+	}
+	let config = {};
+	// globalThis.
+	initialize();
+	renderVis();
+	function initialize() {
+		// globalThis.initialize = () => {
+			console.log("init")
+		config.attention = params['attention'];
+		config.filter = params['default_filter'];
+		config.rootDivId = params['root_div_id'];
+		config.nLayers = config.attention[config.filter]['attn'].length;
+		config.nHeads = config.attention[config.filter]['attn'][0].length;
+		config.layers = params['include_layers']
+		if (params['heads']) {
+			config.headVis = new Array(config.nHeads).fill(false);
+			params['heads'].forEach(x => config.headVis[x] = true);
+		} else {
+			config.headVis = new Array(config.nHeads).fill(true);
+		}
+		config.initialTextLength = config.attention[config.filter].right_text.length;
+		config.layer_seq = (params['layer'] == null ? 0 : config.layers.findIndex(layer => params['layer'] === layer));
+		config.layer = config.layers[config.layer_seq]
+		// '#' + temp1.root_div_id+ ' #layer'
+		$('#' + config.rootDivId+ ' #layer').empty();
+		let layerEl = $('#' + config.rootDivId+ ' #layer');
+		console.log(layerEl)
+		for (const layer of config.layers) {
+			layerEl.append($("<option />").val(layer).text(layer));
+		}
+		layerEl.val(config.layer).change();
+		layerEl.on('change', function (e) {
+			config.layer = +e.currentTarget.value;
+			config.layer_seq = config.layers.findIndex(layer => config.layer === layer);
+			renderVis();
+		});
+		$('#'+config.rootDivId+' #filter').on('change', function (e) {
+		// $(`#${config.rootDivId} #filter`).on('change', function (e) {
+			config.filter = e.currentTarget.value;
+			renderVis();
+		});
+	}
+	function renderVis() {
+		// Load parameters
+		const attnData = config.attention[config.filter];
+		const leftText = attnData.left_text;
+		const rightText = attnData.right_text;
+		// Select attention for given layer
+		const layerAttention = attnData.attn[config.layer_seq];
+		// Clear vis
+		$('#'+config.rootDivId+' #vis').empty();
+		// Determine size of visualization
+		const height = Math.max(leftText.length, rightText.length) * BOXHEIGHT + TEXT_TOP;
+		const svg = d3.select('#'+ config.rootDivId +' #vis')
+			.append('svg')
+			.attr("width", "100%")
+			.attr("height", height + "px");
+		// Display tokens on left and right side of visualization
+		renderText(svg, leftText, true, layerAttention, 0);
+		renderText(svg, rightText, false, layerAttention, MATRIX_WIDTH + BOXWIDTH);
+		// Render attention arcs
+		renderAttention(svg, layerAttention);
+		// Draw squares at top of visualization, one for each head
+		drawCheckboxes(0, svg, layerAttention);
+	}
+	function renderText(svg, text, isLeft, attention, leftPos) {
+		const textContainer = svg.append("svg:g")
+			.attr("id", isLeft ? "left" : "right");
+		// Add attention highlights superimposed over words
+		textContainer.append("g")
+			.classed("attentionBoxes", true)
+			.selectAll("g")
+			.data(attention)
+			.enter()
+			.append("g")
+			.attr("head-index", (d, i) => i)
+			.selectAll("rect")
+			.data(d => isLeft ? d : transpose(d)) // if right text, transpose attention to get right-to-left weights
+			.enter()
+			.append("rect")
+			.attr("x", function () {
+				var headIndex = +this.parentNode.getAttribute("head-index");
+				return leftPos + boxOffsets(headIndex);
+			})
+			.attr("y", (+1) * BOXHEIGHT)
+			.attr("width", BOXWIDTH / activeHeads())
+			.attr("height", BOXHEIGHT)
+			.attr("fill", function () {
+				return headColors(+this.parentNode.getAttribute("head-index"))
+			})
+			.style("opacity", 0.0);
+		const tokenContainer = textContainer.append("g").selectAll("g")
+			.data(text)
+			.enter()
+			.append("g");
+		// Add gray background that appears when hovering over text
+		tokenContainer.append("rect")
+			.classed("background", true)
+			.style("opacity", 0.0)
+			.attr("fill", "lightgray")
+			.attr("x", leftPos)
+			.attr("y", (d, i) => TEXT_TOP + i * BOXHEIGHT)
+			.attr("width", BOXWIDTH)
+			.attr("height", BOXHEIGHT);
+		// Add token text
+		const textEl = tokenContainer.append("text")
+			.text(d => d)
+			.attr("font-size", TEXT_SIZE + "px")
+			.style("cursor", "default")
+			.style("-webkit-user-select", "none")
+			.attr("x", leftPos)
+			.attr("y", (d, i) => TEXT_TOP + i * BOXHEIGHT);
+		if (isLeft) {
+			textEl.style("text-anchor", "end")
+				.attr("dx", BOXWIDTH - 0.5 * TEXT_SIZE)
+				.attr("dy", TEXT_SIZE);
+		} else {
+			textEl.style("text-anchor", "start")
+				.attr("dx", +0.5 * TEXT_SIZE)
+				.attr("dy", TEXT_SIZE);
+		}
+		tokenContainer.on("mouseover", function (d, index) {
+			// Show gray background for moused-over token
+			textContainer.selectAll(".background")
+				.style("opacity", (d, i) => i === index ? 1.0 : 0.0)
+			// Reset visibility attribute for any previously highlighted attention arcs
+			svg.select("#attention")
+				.selectAll("line[visibility='visible']")
+				.attr("visibility", null)
+			// Hide group containing attention arcs
+			svg.select("#attention").attr("visibility", "hidden");
+			// Set to visible appropriate attention arcs to be highlighted
+			if (isLeft) {
+				svg.select("#attention").selectAll("line[left-token-index='" + index + "']").attr("visibility", "visible");
+			} else {
+				svg.select("#attention").selectAll("line[right-token-index='" + index + "']").attr("visibility", "visible");
+			}
+			// Update color boxes superimposed over tokens
+			const id = isLeft ? "right" : "left";
+			const leftPos = isLeft ? MATRIX_WIDTH + BOXWIDTH : 0;
+			svg.select("#" + id)
+				.selectAll(".attentionBoxes")
+				.selectAll("g")
+				.attr("head-index", (d, i) => i)
+				.selectAll("rect")
+				.attr("x", function () {
+					const headIndex = +this.parentNode.getAttribute("head-index");
+					return leftPos + boxOffsets(headIndex);
+				})
+				.attr("y", (d, i) => TEXT_TOP + i * BOXHEIGHT)
+				.attr("width", BOXWIDTH / activeHeads())
+				.attr("height", BOXHEIGHT)
+				.style("opacity", function (d) {
+					const headIndex = +this.parentNode.getAttribute("head-index");
+					if (config.headVis[headIndex])
+						if (d) {
+							return d[index];
+						} else {
+							return 0.0;
+						}
+					else
+						return 0.0;
+				});
+		});
+		textContainer.on("mouseleave", function () {
+			// Unhighlight selected token
+			d3.select(this).selectAll(".background")
+				.style("opacity", 0.0);
+			// Reset visibility attributes for previously selected lines
+			svg.select("#attention")
+				.selectAll("line[visibility='visible']")
+				.attr("visibility", null) ;
+			svg.select("#attention").attr("visibility", "visible");
+			// Reset highlights superimposed over tokens
+			svg.selectAll(".attentionBoxes")
+				.selectAll("g")
+				.selectAll("rect")
+				.style("opacity", 0.0);
+		});
+	}
+	function renderAttention(svg, attention) {
+		// Remove previous dom elements
+		svg.select("#attention").remove();
+		// Add new elements
+		svg.append("g")
+			.attr("id", "attention") // Container for all attention arcs
+			.selectAll(".headAttention")
+			.data(attention)
+			.enter()
+			.append("g")
+			.classed("headAttention", true) // Group attention arcs by head
+			.attr("head-index", (d, i) => i)
+			.selectAll(".tokenAttention")
+			.data(d => d)
+			.enter()
+			.append("g")
+			.classed("tokenAttention", true) // Group attention arcs by left token
+			.attr("left-token-index", (d, i) => i)
+			.selectAll("line")
+			.data(d => d)
+			.enter()
+			.append("line")
+			.attr("x1", BOXWIDTH)
+			.attr("y1", function () {
+				const leftTokenIndex = +this.parentNode.getAttribute("left-token-index")
+				return TEXT_TOP + leftTokenIndex * BOXHEIGHT + (BOXHEIGHT / 2)
+			})
+			.attr("x2", BOXWIDTH + MATRIX_WIDTH)
+			.attr("y2", (d, rightTokenIndex) => TEXT_TOP + rightTokenIndex * BOXHEIGHT + (BOXHEIGHT / 2))
+			.attr("stroke-width", 2)
+			.attr("stroke", function () {
+				const headIndex = +this.parentNode.parentNode.getAttribute("head-index");
+				return headColors(headIndex)
+			})
+			.attr("left-token-index", function () {
+				return +this.parentNode.getAttribute("left-token-index")
+			})
+			.attr("right-token-index", (d, i) => i)
+		;
+		updateAttention(svg)
+	}
+	function updateAttention(svg) {
+		svg.select("#attention")
+			.selectAll("line")
+			.attr("stroke-opacity", function (d) {
+				const headIndex = +this.parentNode.parentNode.getAttribute("head-index");
+				// If head is selected
+				if (config.headVis[headIndex]) {
+					// Set opacity to attention weight divided by number of active heads
+					return d / activeHeads()
+				} else {
+					return 0.0;
+				}
+			})
+	}
+	function boxOffsets(i) {
+		const numHeadsAbove = config.headVis.reduce(
+			function (acc, val, cur) {
+				return val && cur < i ? acc + 1 : acc;
+			}, 0);
+		return numHeadsAbove * (BOXWIDTH / activeHeads());
+	}
+	function activeHeads() {
+		return config.headVis.reduce(function (acc, val) {
+			return val ? acc + 1 : acc;
+		}, 0);
+	}
+	function drawCheckboxes(top, svg) {
+		const checkboxContainer = svg.append("g");
+		const checkbox = checkboxContainer.selectAll("rect")
+			.data(config.headVis)
+			.enter()
+			.append("rect")
+			.attr("fill", (d, i) => headColors(i))
+			.attr("x", (d, i) => i * CHECKBOX_SIZE)
+			.attr("y", top)
+			.attr("width", CHECKBOX_SIZE)
+			.attr("height", CHECKBOX_SIZE);
+		function updateCheckboxes() {
+			checkboxContainer.selectAll("rect")
+				.data(config.headVis)
+				.attr("fill", (d, i) => d ? headColors(i): lighten(headColors(i)));
+		}
+		updateCheckboxes();
+		checkbox.on("click", function (d, i) {
+			if (config.headVis[i] && activeHeads() === 1) return;
+			config.headVis[i] = !config.headVis[i];
+			updateCheckboxes();
+			updateAttention(svg);
+		});
+		checkbox.on("dblclick", function (d, i) {
+			// If we double click on the only active head then reset
+			if (config.headVis[i] && activeHeads() === 1) {
+				config.headVis = new Array(config.nHeads).fill(true);
+			} else {
+				config.headVis = new Array(config.nHeads).fill(false);
+				config.headVis[i] = true;
+			}
+			updateCheckboxes();
+			updateAttention(svg);
+		});
+	}
+	function lighten(color) {
+		const c = d3.hsl(color);
+		const increment = (1 - c.l) * 0.6;
+		c.l += increment;
+		c.s -= increment;
+		return c;
+	}
+	function transpose(mat) {
+		return mat[0].map(function (col, i) {
+			return mat.map(function (row) {
+				return row[i];
+			});
+		});
+	}
+	}
+	// );
+}

requirements.txt ADDED Viewed

	@@ -0,0 +1,2 @@


1	+ inseq
2	+ bertviz