FreshBench / gradio_samples /bertviz /plotsjs_bertviz.js
jijivski
todo added and bertviz base success
124b675
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];
});
});
}
}
// );
}