scope esm2 config

Author: aditya0by0

chebai_proteins/preprocessing/datasets/scope/scope.py

@@ -24,7 +24,7 @@
 import torch
 from Bio import SeqIO
 from chebai.preprocessing.datasets.base import _DynamicDataset
-from chebai.preprocessing.reader import ProteinDataReader
+from chebai.preprocessing.reader import ESM2EmbeddingReader, ProteinDataReader
 
 
 class _SCOPeDataExtractor(_DynamicDataset, ABC):
@@ -952,6 +952,10 @@ class SCOPeOverPartial2000(_SCOPeOverXPartial):
     THRESHOLD: int = 2000
 
 
+class SCOPEOver50ESM(SCOPeOver50):
+    READER = ESM2EmbeddingReader
+
+
 if __name__ == "__main__":
     scope = SCOPeOver50(scope_version="2.08")
 

configs/data/scope/scope50_esm.yml

@@ -0,0 +1,6 @@
+class_path: chebai.preprocessing.datasets.scope.scope.SCOPeOver50ESM
+init_args:
+  scope_version: "2.08"
+  reader_kwargs: {
+    truncation_length: 1000
+  }