add descriptions for 2 pretrained models

Author: lucidrains

enformer_pytorch/config.yml

@@ -1,5 +1,7 @@
 preview:
     id: 1-5nXoOCcRQxFULV4gac9GikY9j3iO3p3
+    description: >
+        Trained on sequences of 131,072 basepairs, target length 896 on v3-64 TPUs for 2 and a half days without augmentations and poisson loss
     params:
         dim: 1536
         depth: 11
@@ -10,6 +12,8 @@ preview:
 
 corr_coef_obj:
     id: 18gMf2gIbU5LV09_2w8yMjEesVzwu4sGz
+    description: >
+        Trained on sequences o 131,072 baespairs, target length 896 on v3-64 TPUs for 3 days with sequence augmentations and pearson correlation objective
     params:
         dim: 1536
         depth: 11

enformer_pytorch/model_loader.py

@@ -36,6 +36,7 @@ def load_pretrained_model(
     slug,
     force = False,
     model = None,
+    verbose = True,
     **kwargs
 ):
     if slug not in CONFIG:
@@ -67,5 +68,8 @@ def load_pretrained_model(
 
     model.load_state_dict(torch.load(str(save_path)))
 
-    print(f'loaded {slug} successfully')
+    if verbose and 'description' in config:
+        print(f"\n{config['description']}\n")
+
+    print(f'Enformer model "{slug}" loaded successfully')
     return model

setup.py

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