Added numbers to BertViz visualization, easier to map attention head numbers

plans.md

@@ -1,5 +1,4 @@
 To Do:
-- generating with 1 output token doesn't show the "selected" beam, so users can't see the output from the model
 - add head index numbers to color swatches in BertViz visualization graph
 - clicking send in the chatbot doesn't immediately show the purple text bubble from the user and the typing animation, there are a couple seconds of nothing. It should immediately populate the text bubble and start the typing animation
 
@@ -13,4 +12,5 @@ Done:
 - shorter, concise responses in system prompt
 - add video links to glossary
     - three blue one brown
-- specs on what each attention head does
+- specs on what each attention head does
+- generating with 1 output token doesn't show the "selected" beam, so users can't see the output from the model

utils/model_patterns.py

@@ -1383,6 +1383,31 @@ def generate_bertviz_html(activation_data: Dict[str, Any], layer_index: int, vie
                 'headColors = d3.scale.category10();',
                 _patch
             )
+
+            # Inject head-index labels inside the checkbox swatches.
+            # Target the first `updateCheckboxes();` call inside drawCheckboxes
+            # and prepend D3 code that appends <text> elements over each rect.
+            _label_js = (
+                'checkboxContainer.selectAll("text")\n'
+                '            .data(config.headVis)\n'
+                '            .enter()\n'
+                '            .append("text")\n'
+                '            .text((d, i) => i)\n'
+                '            .attr("x", (d, i) => i * CHECKBOX_SIZE + CHECKBOX_SIZE / 2)\n'
+                '            .attr("y", top + CHECKBOX_SIZE / 2)\n'
+                '            .attr("text-anchor", "middle")\n'
+                '            .attr("dominant-baseline", "central")\n'
+                '            .attr("font-size", "10px")\n'
+                '            .attr("font-weight", "bold")\n'
+                '            .attr("fill", "white")\n'
+                '            .attr("pointer-events", "none");\n'
+                '        updateCheckboxes();'
+            )
+            html_str = html_str.replace(
+                'updateCheckboxes();\n\n        checkbox.on',
+                _label_js + '\n\n        checkbox.on',
+                1,  # replace only the first occurrence
+            )
             return html_str
             
     except Exception as e: