add imports

.gitignore

@@ -152,3 +152,5 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+.vscode

tests/integration_tests/ComplexInterventionWithGPT2TestCase.py

@@ -1,5 +1,6 @@
 import unittest
 from ..utils import *
+from transformers import GPT2Config
 
 
 class ComplexInterventionWithGPT2TestCase(unittest.TestCase):

tests/integration_tests/InterventionWithGPT2TestCase.py

@@ -1,5 +1,6 @@
 import unittest
 from ..utils import *
+from transformers import GPT2Config
 
 
 class InterventionWithGPT2TestCase(unittest.TestCase):

tests/unit_tests/IntervenableConfigUnitTestCase.py

@@ -1,5 +1,6 @@
 import unittest
 from ..utils import *
+from transformers import GPT2Config
 
 
 class IntervenableConfigUnitTestCase(unittest.TestCase):

tests/unit_tests/InterventionUtilsTestCase.py

@@ -4,6 +4,7 @@
 from pyvene.models.intervention_utils import _do_intervention_by_swap
 from pyvene.models.interventions import VanillaIntervention
 from pyvene.models.interventions import CollectIntervention
+from transformers import GPT2Config
 
 
 class InterventionUtilsTestCase(unittest.TestCase):

tests/unit_tests/ModelUtilsTestCase.py

@@ -2,6 +2,7 @@
 from ..utils import *
 from pyvene.models.modeling_utils import *
 from pprint import pprint, pformat
+from transformers import GPT2Config
 
 
 class ModelUtilsTestCase(unittest.TestCase):