Merge pull request #24 from JonathanDZiegler/fix_bug_in_TRRMSADataset

Fix bug in trrmsa dataset
microsoft · Nov 13, 2023 · 32e3bd8 · 32e3bd8
2 parents 418223d + cf61dd1
commit 32e3bd8
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/evodiff/data.py b/evodiff/data.py
@@ -253,7 +253,7 @@ def __getitem__(self, idx):
             slice_start = 0
             seq_len = msa_seq_len
 
-        sliced_msa = msa[:, slice_start: slice_start + self.max_seq_len]
+        sliced_msa = msa[:, slice_start: slice_start + seq_len]
         anchor_seq = sliced_msa[0]  # This is the query sequence in MSA
 
         sliced_msa = [list(seq) for seq in sliced_msa if (list(set(seq)) != [self.tokenizer.alphabet.index(GAP)])]
@@ -271,7 +271,7 @@ def __getitem__(self, idx):
                 anchor_seq = np.expand_dims(anchor_seq, axis=0)
                 output = np.concatenate((anchor_seq, sliced_msa[random_idx]), axis=0)
             elif self.selection_type == 'non-random':
-                output = sliced_msa[:64]
+                output = sliced_msa[:self.n_sequences]
             elif self.selection_type == "MaxHamming":
                 output = [list(anchor_seq)]
                 msa_subset = sliced_msa[1:]