able to return nucleotide indices from dataset

Author: lucidrains

README.md

@@ -276,6 +276,7 @@ ds = GenomeIntervalDataset(
     bed_file = './sequences.bed',                       # bed file
     fasta_file = './hg38.ml.fa',                        # path to fasta file
     filter_df_fn = lambda df: df[df.type == 'train'],   # filter dataframe function
+    return_seq_indices = True,                          # return nucleotide indices (ACGTN) or one hot encodings
     context_length = 196_608,
     # this can be longer than the interval designated in the .bed file,
     # in which case it will take care of lengthening the interval on either sides
@@ -290,7 +291,8 @@ model = Enformer(
     target_length = 896,
 )
 
-pred = model(ds[0], head = 'human') # (896, 5313)
+seq = ds[0] # (196608,)
+pred = model(seq, head = 'human') # (896, 5313)
 ```
 
 ## Appreciation

enformer_pytorch/data.py

@@ -38,7 +38,7 @@ def __init__(
         bed_file,
         fasta_file,
         context_length = None,
-        shift_augmentation_range = None,
+        return_seq_indices = False,
         filter_df_fn = identity
     ):
         super().__init__()
@@ -54,7 +54,7 @@ def __init__(
         self.df = df
         self.seqs = Fasta(str(fasta_file))
         self.context_length = context_length
-        self.shift_augmentation = shift_augmentation
+        self.return_seq_indices = return_seq_indices
 
     def __len__(self):
         return len(self.df)
@@ -88,5 +88,9 @@ def __getitem__(self, ind):
 
         seq = ('.' * left_padding) + str(chromosome[start:end]) + ('.' * right_padding)
         seq_indices = str_to_seq_indices(seq)
+
+        if self.return_seq_indices:
+            return seq_indices.squeeze(0)
+
         seq_onehot = seq_indices_to_one_hot(seq_indices)
         return seq_onehot.squeeze(0)

setup.py

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