New text neuron activations args: show_selectors and dimension_default

.gitignore

@@ -11,6 +11,7 @@ __pycache__
 venv
 *-checkpoint.ipynb
 requirements.txt
+python/circuitsvis/.vscode
 
 # Node
 node_modules/

python/circuitsvis/activations.py

@@ -13,6 +13,9 @@ def text_neuron_activations(
     second_dimension_name: Optional[str] = "Neuron",
     first_dimension_labels: Optional[List[str]] = None,
     second_dimension_labels: Optional[List[str]] = None,
+    first_dimension_default: Optional[int] = 0,
+    second_dimension_default: Optional[int] = 0,
+    show_selectors: Optional[bool] = True,
 ) -> RenderedHTML:
     """Show activations (colored by intensity) for each token in a text or set
     of texts.
@@ -54,4 +57,7 @@ def text_neuron_activations(
         secondDimensionName=second_dimension_name,
         firstDimensionLabels=first_dimension_labels,
         secondDimensionLabels=second_dimension_labels,
+        firstDimensionDefault=first_dimension_default,
+        secondDimensionDefault=second_dimension_default,
+        showSelectors=show_selectors,
     )

python/circuitsvis/tests/snapshots/snap_test_activations.py

@@ -7,22 +7,10 @@
 
 snapshots = Snapshot()
 
-snapshots['TestTextNeuronActivations.test_multi_matches_snapshot 1'] = '''<div id="circuits-vis-mock" style="margin: 15px 0;"/>
-    <script crossorigin type="module">
-    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";
-    render(
-      "circuits-vis-mock",
-      TextNeuronActivations,
-      {"tokens": [["a", "b"], ["c", "d", "e"]], "activations": [[[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]], [[0, 1, 1], [1, 1, 1]]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron"}
-    )
-    </script>'''
+snapshots[
+    "TestTextNeuronActivations.test_multi_matches_snapshot 1"
+] = """<div id="circuits-vis-mock" style="margin: 15px 0;"/>\n    <script crossorigin type="module">\n    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";\n    render(\n      "circuits-vis-mock",\n      TextNeuronActivations,\n      {"tokens": [["a", "b"], ["c", "d", "e"]], "activations": [[[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]], [[0, 1, 1], [1, 1, 1]]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron", "firstDimensionDefault": 0, "secondDimensionDefault": 0, "showSelectors": true}\n    )\n    </script>"""
 
-snapshots['TestTextNeuronActivations.test_single_matches_snapshot 1'] = '''<div id="circuits-vis-mock" style="margin: 15px 0;"/>
-    <script crossorigin type="module">
-    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";
-    render(
-      "circuits-vis-mock",
-      TextNeuronActivations,
-      {"tokens": ["a", "b"], "activations": [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron"}
-    )
-    </script>'''
+snapshots[
+    "TestTextNeuronActivations.test_single_matches_snapshot 1"
+] = """<div id="circuits-vis-mock" style="margin: 15px 0;"/>\n    <script crossorigin type="module">\n    import { render, TextNeuronActivations } from "https://unpkg.com/circuitsvis@0.0.0/dist/cdn/esm.js";\n    render(\n      "circuits-vis-mock",\n      TextNeuronActivations,\n      {"tokens": ["a", "b"], "activations": [[[0, 1, 0], [0, 1, 1]], [[0, 1, 1], [1, 1, 1]]], "firstDimensionName": "Layer", "secondDimensionName": "Neuron", "firstDimensionDefault": 0, "secondDimensionDefault": 0, "showSelectors": true}\n    )\n    </script>"""

python/circuitsvis/topk_samples.py

@@ -1,12 +1,13 @@
 """Activations visualizations"""
-from typing import List, Optional
+import torch
+from typing import List, Optional, Union
 
 from circuitsvis.utils.render import RenderedHTML, render
 
 
 def topk_samples(
     tokens: List[List[List[List[str]]]],
-    activations: List[List[List[List[float]]]],
+    activations: Union[List[List[List[List[float]]]], torch.Tensor],
     zeroth_dimension_name: Optional[str] = "Layer",
     first_dimension_name: Optional[str] = "Neuron",
     zeroth_dimension_labels: Optional[List[str]] = None,

react/src/activations/TextNeuronActivations.tsx

@@ -36,7 +36,10 @@ export function TextNeuronActivations({
   firstDimensionName = "Layer",
   secondDimensionName = "Neuron",
   firstDimensionLabels,
-  secondDimensionLabels
+  secondDimensionLabels,
+  firstDimensionDefault = 0,
+  secondDimensionDefault = 0,
+  showSelectors = true
 }: TextNeuronActivationsProps) {
   // If there is only one sample (i.e. if tokens is an array of strings), cast tokens and activations to an array with
   // a single element
@@ -68,8 +71,10 @@ export function TextNeuronActivations({
   const [sampleNumbers, setSampleNumbers] = useState<number[]>([
     ...Array(samplesPerPage).keys()
   ]);
-  const [layerNumber, setLayerNumber] = useState<number>(0);
-  const [neuronNumber, setNeuronNumber] = useState<number>(0);
+  const [layerNumber, setLayerNumber] = useState<number>(firstDimensionDefault);
+  const [neuronNumber, setNeuronNumber] = useState<number>(
+    secondDimensionDefault
+  );
 
   useEffect(() => {
     // When the user changes the samplesPerPage, update the sampleNumbers
@@ -96,77 +101,79 @@ export function TextNeuronActivations({
 
   return (
     <Container fluid>
-      <Row>
-        <Col>
-          <Row style={selectRowStyle}>
-            <Col>
-              <label htmlFor="layer-selector" style={{ marginRight: 15 }}>
-                {firstDimensionName}:
-              </label>
-              <NumberSelector
-                id="layer-selector"
-                largestNumber={numberOfLayers! - 1}
-                currentValue={layerNumber}
-                setCurrentValue={setLayerNumber}
-                labels={firstDimensionLabels}
-              />
-            </Col>
-          </Row>
-          <Row style={selectRowStyle}>
-            <Col>
-              <label htmlFor="neuron-selector" style={{ marginRight: 15 }}>
-                {secondDimensionName}:
-              </label>
-              <NumberSelector
-                id="neuron-selector"
-                largestNumber={numberOfNeurons! - 1}
-                currentValue={neuronNumber}
-                setCurrentValue={setNeuronNumber}
-                labels={secondDimensionLabels}
-              />
-            </Col>
-          </Row>
-          {/* Only show the sample selector if there is more than one sample */}
-          {numberOfSamples > 1 && (
+      {showSelectors && (
+        <Row>
+          <Col>
             <Row style={selectRowStyle}>
               <Col>
-                <label htmlFor="sample-selector" style={{ marginRight: 15 }}>
-                  Samples:
+                <label htmlFor="layer-selector" style={{ marginRight: 15 }}>
+                  {firstDimensionName}:
                 </label>
-                <RangeSelector
-                  id="sample-selector"
-                  largestNumber={numberOfSamples - 1}
-                  currentRangeArr={sampleNumbers}
-                  setCurrentValue={setSampleNumbers}
-                  numValsInRange={samplesPerPage}
+                <NumberSelector
+                  id="layer-selector"
+                  largestNumber={numberOfLayers! - 1}
+                  currentValue={layerNumber}
+                  setCurrentValue={setLayerNumber}
+                  labels={firstDimensionLabels}
                 />
               </Col>
             </Row>
-          )}
-        </Col>
-        <Col>
-          {/* Only show the sample per page selector if there is more than one sample */}
-          {numberOfSamples > 1 && (
             <Row style={selectRowStyle}>
               <Col>
-                <label
-                  htmlFor="samples-per-page-selector"
-                  style={{ marginRight: 15 }}
-                >
-                  Samples per page:
+                <label htmlFor="neuron-selector" style={{ marginRight: 15 }}>
+                  {secondDimensionName}:
                 </label>
                 <NumberSelector
-                  id="samples-per-page-selector"
-                  smallestNumber={1}
-                  largestNumber={numberOfSamples}
-                  currentValue={samplesPerPage}
-                  setCurrentValue={setSamplesPerPage}
+                  id="neuron-selector"
+                  largestNumber={numberOfNeurons! - 1}
+                  currentValue={neuronNumber}
+                  setCurrentValue={setNeuronNumber}
+                  labels={secondDimensionLabels}
                 />
               </Col>
             </Row>
-          )}
-        </Col>
-      </Row>
+            {/* Only show the sample selector if there is more than one sample */}
+            {numberOfSamples > 1 && (
+              <Row style={selectRowStyle}>
+                <Col>
+                  <label htmlFor="sample-selector" style={{ marginRight: 15 }}>
+                    Samples:
+                  </label>
+                  <RangeSelector
+                    id="sample-selector"
+                    largestNumber={numberOfSamples - 1}
+                    currentRangeArr={sampleNumbers}
+                    setCurrentValue={setSampleNumbers}
+                    numValsInRange={samplesPerPage}
+                  />
+                </Col>
+              </Row>
+            )}
+          </Col>
+          <Col>
+            {/* Only show the sample per page selector if there is more than one sample */}
+            {numberOfSamples > 1 && (
+              <Row style={selectRowStyle}>
+                <Col>
+                  <label
+                    htmlFor="samples-per-page-selector"
+                    style={{ marginRight: 15 }}
+                  >
+                    Samples per page:
+                  </label>
+                  <NumberSelector
+                    id="samples-per-page-selector"
+                    smallestNumber={1}
+                    largestNumber={numberOfSamples}
+                    currentValue={samplesPerPage}
+                    setCurrentValue={setSamplesPerPage}
+                  />
+                </Col>
+              </Row>
+            )}
+          </Col>
+        </Row>
+      )}
       <Row>
         <Col>
           <SampleItems
@@ -218,4 +225,19 @@ export interface TextNeuronActivationsProps {
    * Labels for the second dimension
    */
   secondDimensionLabels?: string[];
+
+  /**
+   * Default index for the first dimension
+   */
+  firstDimensionDefault?: number;
+
+  /**
+   * Default index for the second dimension
+   */
+  secondDimensionDefault?: number;
+
+  /**
+   * Whether to show the selector dropdowns
+   */
+  showSelectors?: boolean;
 }