default max len of go, scope should be same

- needed for pretraining (max_position_embeddings)

Author: aditya0by0

chebai_proteins/preprocessing/datasets/deepGO/go_uniprot.py

@@ -102,17 +102,18 @@ class _GOUniProtDataExtractor(_DynamicDataset, ABC):
 
     # Gene Ontology (GO) has three major branches, one for biological processes (BP), molecular functions (MF) and
     # cellular components (CC). The value "all" will take data related to all three branches into account.
+    # TODO: should we be really allowing all branches for single dataset?
     _ALL_GO_BRANCHES: str = "all"
     _GO_BRANCH_NAMESPACE: Dict[str, str] = {
         "BP": "biological_process",
         "MF": "molecular_function",
         "CC": "cellular_component",
     }
 
-    def __init__(self, **kwargs):
-        self.go_branch: str = self._get_go_branch(**kwargs)
+    def __init__(self, go_branch: str, max_sequence_len: int = 1002, **kwargs):
+        self.go_branch: str = self._get_go_branch(go_branch)
 
-        self.max_sequence_length: int = int(kwargs.get("max_sequence_length", 1002))
+        self.max_sequence_length: int = int(max_sequence_len)
         assert (
             self.max_sequence_length >= 1
         ), "Max sequence length should be greater than or equal to 1."
@@ -126,7 +127,7 @@ def __init__(self, **kwargs):
             )
 
     @classmethod
-    def _get_go_branch(cls, **kwargs) -> str:
+    def _get_go_branch(cls, go_branch_value: str, **kwargs) -> str:
         """
         Retrieves the Gene Ontology (GO) branch based on provided keyword arguments.
         This method checks if a valid GO branch value is provided in the keyword arguments.
@@ -141,7 +142,6 @@ def _get_go_branch(cls, **kwargs) -> str:
             ValueError: If the provided 'go_branch' value is not in the allowed list of values.
         """
 
-        go_branch_value = kwargs.get("go_branch", cls._ALL_GO_BRANCHES)
         allowed_values = list(cls._GO_BRANCH_NAMESPACE.keys()) + [cls._ALL_GO_BRANCHES]
         if go_branch_value not in allowed_values:
             raise ValueError(

chebai_proteins/preprocessing/datasets/scope/scope.py

@@ -72,12 +72,12 @@ def __init__(
         self,
         scope_version: str,
         scope_version_train: Optional[str] = None,
-        max_sequence_len: int = 1000,
+        max_sequence_len: int = 1002,
         **kwargs,
     ):
         self.scope_version: str = scope_version
         self.scope_version_train: str = scope_version_train
-        self.max_sequence_len: int = max_sequence_len
+        self.max_sequence_len: int = int(max_sequence_len)
 
         super(_SCOPeDataExtractor, self).__init__(**kwargs)
 
@@ -224,7 +224,6 @@ def add_sequence_nodes_edges(chain_sequence, px_sun_id):
         # Step 1: Build the graph structure and store node attributes
         for row in df_scope.itertuples(index=False):
             if row.level == "px":
-
                 pdb_id, chain_id = row.sid[1:5], row.sid[5]
 
                 if pdb_id not in pdb_id_set or chain_id == "_":
@@ -546,7 +545,6 @@ def _parse_pdb_sequence_file(self) -> pd.DataFrame:
         for record in SeqIO.parse(
             os.path.join(self.scope_root_dir, self.raw_file_names_dict["PDB"]), "fasta"
         ):
-
             if not record.seq or len(record.seq) > self.max_sequence_len:
                 continue
 
@@ -934,7 +932,6 @@ class SCOPeOver2000(_SCOPeOverX):
 
 
 class SCOPeOver50(_SCOPeOverX):
-
     THRESHOLD = 50
 
 

configs/model/electra.yml

@@ -4,7 +4,7 @@ init_args:
     lr: 1e-3
   config:
     vocab_size: 31
-    max_position_embeddings: 1000
+    max_position_embeddings: 1002
     num_attention_heads: 8
     num_hidden_layers: 6
     type_vocab_size: 1