Attempted to add show_selectors arg

TransformerLensOrg · Sep 19, 2023 · 113d10d · 113d10d
1 parent 19e9281
commit 113d10d
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 59 deletions.
diff --git a/python/circuitsvis/activations.py b/python/circuitsvis/activations.py
@@ -15,6 +15,7 @@ def text_neuron_activations(
     second_dimension_labels: Optional[List[str]] = None,
     first_dimension_default: Optional[int] = 0,
     second_dimension_default: Optional[int] = 0,
+    show_selectors: bool = True,
 ) -> RenderedHTML:
     """Show activations (colored by intensity) for each token in a text or set
     of texts.
@@ -58,4 +59,5 @@ def text_neuron_activations(
         secondDimensionLabels=second_dimension_labels,
         firstDimensionDefault=first_dimension_default,
         secondDimensionDefault=second_dimension_default,
+        showSelectors=show_selectors,
     )
diff --git a/react/src/activations/TextNeuronActivations.tsx b/react/src/activations/TextNeuronActivations.tsx
@@ -38,7 +38,8 @@ export function TextNeuronActivations({
   firstDimensionLabels,
   secondDimensionLabels,
   firstDimensionDefault = 0,
-  secondDimensionDefault = 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
@@ -98,77 +99,62 @@ 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>
+          </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
@@ -231,4 +217,9 @@ export interface TextNeuronActivationsProps {
    */
   secondDimensionDefault?: number;
 
+
+    /**
+   * Whether to show the selector dropdowns
+   */
+  showSelectors?: boolean;
 }