pyvene unit test passing

pyvene/models/gru/modelings_intervenable_gru.py

@@ -100,7 +100,7 @@ def create_gru_classifier(
 ):
     """Creates a GRU model, config, and tokenizer from the given name and revision"""
     from transformers import AutoTokenizer
-    from models.gru.modelings_gru import GRUForClassification
+    from pyvene.models.gru.modelings_gru import GRUForClassification
 
     tokenizer = None
     if tokenizer_name is not None:

pyvene/models/mlp/modelings_intervenable_mlp.py

@@ -40,7 +40,7 @@ def create_mlp_classifier(
 ):
     """Creates a MLP model, config, and tokenizer from the given name and revision"""
     from transformers import AutoTokenizer
-    from models.mlp.modelings_mlp import MLPForClassification
+    from pyvene.models.mlp.modelings_mlp import MLPForClassification
 
     tokenizer = None
     if tokenizer_name is not None:

tests/utils.py

@@ -8,20 +8,20 @@
 import numpy as np
 from transformers import GPT2Config
 
-from models.basic_utils import embed_to_distrib, top_vals, format_token
-from models.configuration_intervenable_model import (
+from pyvene.models.basic_utils import embed_to_distrib, top_vals, format_token
+from pyvene.models.configuration_intervenable_model import (
     IntervenableRepresentationConfig,
     IntervenableConfig,
 )
-from models.intervenable_base import IntervenableModel
-from models.interventions import (
+from pyvene.models.intervenable_base import IntervenableModel
+from pyvene.models.interventions import (
     VanillaIntervention,
     RotatedSpaceIntervention,
     LowRankRotatedSpaceIntervention,
 )
-from models.mlp.modelings_mlp import MLPConfig
-from models.mlp.modelings_intervenable_mlp import create_mlp_classifier
-from models.gpt2.modelings_intervenable_gpt2 import create_gpt2_lm
+from pyvene.models.mlp.modelings_mlp import MLPConfig
+from pyvene.models.mlp.modelings_intervenable_mlp import create_mlp_classifier
+from pyvene.models.gpt2.modelings_intervenable_gpt2 import create_gpt2_lm
 
 
 ##################