Merge pull request #137 from amirzur2023/main

[Minor] Revert CausalModel to accept input/output functions when generating factual/counterfactual datasets

Author: frankaging

pyvene/data_generators/causal_model.py

@@ -309,10 +309,17 @@ def generate_factual_dataset(
         sampler=None,
         filter=None,
         device="cpu",
+        input_function=None,
+        output_function=None,
         return_tensors=True,
     ):
         if sampler is None:
             sampler = self.sample_input
+        
+        if input_function is None:
+            input_function = self.input_to_tensor
+        if output_function is None:
+            output_function = self.output_to_tensor
 
         examples = []
         while len(examples) < size:
@@ -321,8 +328,8 @@ def generate_factual_dataset(
             if filter is None or filter(input):
                 output = self.run_forward(input)
                 if return_tensors:
-                    example['input_ids'] = self.input_to_tensor(input).to(device)
-                    example['labels'] = self.output_to_tensor(output).to(device)
+                    example['input_ids'] = input_function(input).to(device)
+                    example['labels'] = output_function(output).to(device)
                 else:
                     example['input_ids'] = input
                     example['labels'] = output
@@ -339,8 +346,15 @@ def generate_counterfactual_dataset(
         intervention_sampler=None,
         filter=None,
         device="cpu",
+        input_function=None,
+        output_function=None,
         return_tensors=True,
     ):
+        if input_function is None:
+            input_function = self.input_to_tensor
+        if output_function is None:
+            output_function = self.output_to_tensor
+
         maxlength = len(
             [
                 var

setup.py

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