Update protein_pretraining.py

Author: aditya0by0

chebai/preprocessing/datasets/protein_pretraining.py

@@ -40,8 +40,20 @@ def __init__(self, **kwargs):
         """
         self._go_uniprot_extractor = GOUniProtOver250()
         assert self._go_uniprot_extractor.go_branch == GOUniProtOver250._ALL_GO_BRANCHES
+
+        self.max_sequence_length: int = int(kwargs.get("max_sequence_length", 1002))
+        assert (
+            self.max_sequence_length >= 1
+        ), "Max sequence length should be greater than or equal to 1."
+
         super(_ProteinPretrainingData, self).__init__(**kwargs)
 
+        if self.reader.n_gram is not None:
+            assert self.max_sequence_length >= self.reader.n_gram, (
+                f"max_sequence_length ({self.max_sequence_length}) must be greater than "
+                f"or equal to n_gram ({self.reader.n_gram})."
+            )
+
     # ------------------------------ Phase: Prepare data -----------------------------------
     def prepare_data(self, *args: Any, **kwargs: Any) -> None:
         """
@@ -120,6 +132,10 @@ def _parse_protein_data_for_pretraining(self) -> pd.DataFrame:
                 # Consider protein with only sequence representation
                 continue
 
+            if len(record.sequence) > self.max_sequence_length:
+                # Consider protein with only sequence length not greater than max seq. length
+                continue
+
             if any(aa in AMBIGUOUS_AMINO_ACIDS for aa in record.sequence):
                 # Skip proteins with ambiguous amino acid codes
                 continue
@@ -260,4 +276,4 @@ def _name(self) -> str:
         Returns:
             str: A string identifier, "SwissProteinPretrain", representing the name of this data module.
         """
-        return "SwissProteinPretrain"
+        return f"Swiss_{self.max_sequence_length}"