only allow seq length of 1536 if using tf gamma

Author: lucidrains

enformer_pytorch/modeling_enformer.py

@@ -27,15 +27,6 @@
 DIR = Path(__file__).parents[0]
 TF_GAMMAS = torch.load(str(DIR / "precomputed"/ "tf_gammas.pt"))
 
-def get_tf_gamma(seq_len, device):
-    tf_gammas = TF_GAMMAS.to(device)
-    pad = 1536 - seq_len
-
-    if pad == 0:
-        return tf_gammas
-
-    return tf_gammas[pad:-pad]
-
 # helpers
 
 def exists(val):
@@ -112,10 +103,12 @@ def get_positional_features_gamma(positions, features, seq_len, stddev = None, s
 def get_positional_embed(seq_len, feature_size, device, use_tf_gamma):
     distances = torch.arange(-seq_len + 1, seq_len, device = device)
 
+    assert not use_tf_gamma or seq_len == 1536, 'if using tf gamma, only sequence length of 1536 allowed for now'
+
     feature_functions = [
         get_positional_features_exponential,
         get_positional_features_central_mask,
-        get_positional_features_gamma if not use_tf_gamma else always(get_tf_gamma(seq_len, device))
+        get_positional_features_gamma if not use_tf_gamma else always(TF_GAMMAS.to(device))
     ]
 
     num_components = len(feature_functions) * 2

setup.py

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

test_pretrained.py

@@ -1,7 +1,7 @@
 import torch
 from enformer_pytorch import from_pretrained
 
-enformer = from_pretrained('EleutherAI/enformer-official-rough').cuda()
+enformer = from_pretrained('EleutherAI/enformer-official-rough', use_tf_gamma = False).cuda()
 enformer.eval()
 
 data = torch.load('./data/test-sample.pt')