deepgo: raise error if no classes are selected

Author: aditya0by0

chebai_proteins/preprocessing/datasets/deepGO/go_uniprot.py

@@ -394,6 +394,10 @@ def _graph_to_raw_dataset(self, g: nx.DiGraph) -> pd.DataFrame:
         )
         # Initialize the GO term labels/columns to False
         selected_classes = self.select_classes(g, data_df=data_df)
+        if not selected_classes:
+            raise ValueError(
+                f"No classes selected for given threshold {self.THRESHOLD}"
+            )
         new_label_columns = pd.DataFrame(
             False, index=data_df.index, columns=selected_classes
         )

chebai_proteins/preprocessing/datasets/scope/scope.py

@@ -728,8 +728,6 @@ def _setup_pruned_test_set(
         """
         # TODO: find a more efficient way to do this
         filename_old = "classes.txt"
-        # filename_new = f"classes_v{self.scope_version_train}.txt"
-        # dataset = torch.load(os.path.join(self.processed_dir, "test.pt"))
 
         # Load original classes (from the current SCOPe version - scope_version)
         with open(os.path.join(self.processed_dir_main, filename_old), "r") as file:
@@ -760,7 +758,7 @@ def _setup_pruned_test_set(
                 # set the corresponding label in new_labels to True
                 if mapping[ind] is not None and label:
                     new_labels[mapping[ind]] = label
-            # Update the labels from test instance from scope_version to the new labels, which are compatible to both versions
+            # Update the labels from test instance of scope_version to new labels, which are compatible to both versions
             row["labels"] = new_labels
 
         return df_test_scope_version

chebai_proteins/preprocessing/reader.py

@@ -166,6 +166,7 @@ class ESM2EmbeddingReader(DataReader):
 
     def __init__(
         self,
+        # --------- Default Parameters as per DeepGO2 ------------
         save_model_dir: str = os.path.join("data", "esm2_reader"),
         model_name: str = "esm2_t36_3B_UR50D",
         device: Optional[torch.device] = None,