viz.renderers.AttentionFlowRenderer

class AttentionFlowRenderer

GPU-accelerated attention flow renderer.

Visualizes attention matrices as arcs connecting tokens.

Examples

const renderer = new AttentionFlowRenderer();

// attention: [num_tokens, num_tokens] attention weights
const target = await renderer.render(attention, {
  tokens: ['The', 'cat', 'sat', 'on', 'the', 'mat'],
  highlightToken: 1, // Highlight 'cat'
});