[Minor] Update dependency

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',