Merge pull request #185 from stanfordnlp/zen/dependency_clean

[Minor] Update dependency

Author: frankaging

pyvene/models/intervenable_base.py

@@ -1,5 +1,4 @@
 import json, logging, torch, types
-import nnsight
 import numpy as np
 from collections import OrderedDict
 from typing import List, Optional, Tuple, Union, Dict, Any
@@ -27,6 +26,12 @@
 from transformers.utils import ModelOutput
 from tqdm import tqdm, trange
 
+try:
+    import nnsight
+except:
+    print("nnsight is not detected. Please install via 'pip install nnsight' for nnsight backend.")
+
+
 @dataclass
 class IntervenableModelOutput(ModelOutput):
     """
@@ -226,7 +231,7 @@ def __init__(self, config, model, backend, **kwargs):
         # cached swapped activations (hot)
         self.hot_activations = {}
 
-        self.aux_loss = []
+        self.full_intervention_outputs = []
 
         # temp fields should not be accessed outside
         self._batched_setter_activation_select = {}
@@ -1558,16 +1563,17 @@ def hook_callback(model, args, kwargs, output=None):
                 else:
                     if not isinstance(self.interventions[key][0], types.FunctionType):
                         if intervention.is_source_constant:
-                            intervened_representation = do_intervention(
+                            raw_intervened_representation = do_intervention(
                                 selected_output,
                                 None,
                                 intervention,
                                 subspaces[key_i] if subspaces is not None else None,
                             )
-                            if isinstance(intervened_representation, InterventionOutput):
-                                if intervened_representation.loss is not None:
-                                    self.aux_loss.append(intervened_representation.loss)
-                                intervened_representation = intervened_representation.output
+                            if isinstance(raw_intervened_representation, InterventionOutput):
+                                self.full_intervention_outputs.append(raw_intervened_representation)
+                                intervened_representation = raw_intervened_representation.output
+                            else:
+                                intervened_representation = raw_intervened_representation
                         else:
                             intervened_representation = do_intervention(
                                 selected_output,
@@ -1866,7 +1872,7 @@ def forward(
         if sources is not None and not isinstance(sources, list):
             sources = [sources]
 
-        self.aux_loss.clear()
+        self.full_intervention_outputs.clear()
 
         self._cleanup_states()
 

requirements.txt

@@ -10,4 +10,3 @@ numpy>=1.23.5
 fsspec>=2023.6.0
 accelerate>=0.29.1
 sentencepiece>=0.1.96
-nnsight>=0.1.0

setup.py

@@ -10,7 +10,7 @@
 
 setup(
     name="pyvene",
-    version="0.1.4",
+    version="0.1.5",
     description="Use Activation Intervention to Interpret Causal Mechanism of Model",
     long_description=long_description,
     long_description_content_type='text/markdown',