add esm2 to scope

Author: aditya0by0

chebai_proteins/preprocessing/datasets/deepGO/go_uniprot.py

@@ -120,7 +120,7 @@ def __init__(
         **kwargs,
     ):
         if use_esm2_embeddings:
-            self.READER = dr.ESM2ProteinDataReader
+            self.READER = dr.ESM2EmbeddingReader
 
         self.go_branch: str = self._get_go_branch(go_branch)
 

chebai_proteins/preprocessing/datasets/scope/scope.py

@@ -67,14 +67,19 @@ class _SCOPeDataExtractor(_DynamicDataset, ABC):
         "sp": "species",
         "px": "domain",
     }
+    READER = None
 
     def __init__(
         self,
         scope_version: str,
         scope_version_train: Optional[str] = None,
         max_sequence_len: int = 1002,
+        use_esm2_embedding: bool = False,
         **kwargs,
     ):
+        if use_esm2_embedding:
+            self.READER = ESM2EmbeddingReader
+
         self.scope_version: str = scope_version
         self.scope_version_train: str = scope_version_train
         self.max_sequence_len: int = int(max_sequence_len)
@@ -948,10 +953,6 @@ class SCOPeOverPartial2000(_SCOPeOverXPartial):
     THRESHOLD: int = 2000
 
 
-class SCOPeOver50ESM(SCOPeOver50):
-    READER = ESM2EmbeddingReader
-
-
 if __name__ == "__main__":
     scope = SCOPeOver50(scope_version="2.08")