allow passing model keyword arguments through the parallelize interfa…

…ce (#97)
axonn-ai · Oct 25, 2024 · 6252013 · 6252013
1 parent 660a224
commit 6252013
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/axonn/models/transformers/__init__.py b/axonn/models/transformers/__init__.py
@@ -40,8 +40,8 @@
 
 
 @contextmanager
-def parallelize(model_id):
-    config = AutoConfig.from_pretrained(model_id)
+def parallelize(model_id, model_kwargs={}):
+    config = AutoConfig.from_pretrained(model_id, **model_kwargs)
     architecture = config.architectures[0]
     # config.architectures is a list, not sure what to do
     # if it has multiple elements