allow for passing externally defined enformer model into load_pretrained_model method

Author: lucidrains

README.md

@@ -143,6 +143,18 @@ model = load_pretrained_model('preview', target_length = 128)
 # do your fine-tuning
 ```
 
+You can also define the model externally, and then load the pretrained weights by passing it into `load_pretrained_model`
+
+```python
+from enformer_pytorch import Enformer, load_pretrained_model
+
+enformer = Enformer(dim = 1536, depth = 11, target_length = 128)
+
+load_pretrained_model('preview', model = enformer)
+
+# use enformer
+```
+
 ## Fine-tuning (wip)
 
 This repository will also allow for easy fine-tuning of Enformer.

enformer_pytorch/model_loader.py

@@ -35,7 +35,8 @@ def remove_nones(d):
 def load_pretrained_model(
     slug,
     force = False,
-    target_length = None
+    target_length = None,
+    model = None
 ):
     if slug not in CONFIG:
         print(f'model {slug} not found among available choices: [{", ".join(CONFIG.keys())}]')
@@ -59,7 +60,9 @@ def load_pretrained_model(
     override_params = remove_nones({'target_length': target_length})
     params = {**config['params'], **override_params}
 
-    model = Enformer(**config['params'])
+    if not exists(model):
+        model = Enformer(**config['params'])
+
     model.load_state_dict(torch.load(str(save_path)))
 
     print(f'loaded {slug} successfully')

setup.py

@@ -4,7 +4,7 @@
   name = 'enformer-pytorch',
   packages = find_packages(exclude=[]),
   include_package_data = True,
-  version = '0.1.9',
+  version = '0.1.10',
   license='MIT',
   description = 'Enformer - Pytorch',
   author = 'Phil Wang',