change default masking rate values, and update example notebook

compomics · Jun 21, 2024 · 784e169 · 784e169
1 parent 7b010c3
commit 784e169
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 3 deletions.
diff --git a/molexpress/datasets/encoders.py b/molexpress/datasets/encoders.py
@@ -84,8 +84,8 @@ def collate_fn(
     @staticmethod
     def masked_collate_fn(
         data: list[types.MolecularGraph],
-        node_masking_rate: float = 0.25,
-        edge_masking_rate: float = 0.25,
+        node_masking_rate: float = 0.0,
+        edge_masking_rate: float = 0.0,
     ) -> tuple[types.MolecularGraph, np.ndarray]:
         """
         Merge list of graphs into a single disjoint graph.

diff --git a/notebooks/examples-masking.ipynb b/notebooks/examples-masking.ipynb
@@ -10,6 +10,8 @@
     "import os\n",
     "os.environ[\"KERAS_BACKEND\"] = \"torch\" # Comment out for tensorflow backend\n",
     "\n",
+    "from functools import partial\n",
+    "\n",
     "from molexpress import layers\n",
     "from molexpress.datasets import featurizers\n",
     "from molexpress.datasets import encoders\n",
@@ -117,8 +119,14 @@
     "        \n",
     "torch_dataset = Dataset(x_dummy)\n",
     "\n",
+    "# We should implement the collate functions differently (and probably not as a staticmethod \n",
+    "# of the peptide graph encoder); but for now, we keep it as it is, and we can do a partial \n",
+    "# for additional arguments:\n",
+    "partial_collate_fn = partial(\n",
+    "    peptide_graph_encoder.masked_collate_fn, node_masking_rate=0.25, edge_masking_rate=0.25)\n",
+    "\n",
     "dataset = torch.utils.data.DataLoader(\n",
-    "    torch_dataset, batch_size=2, collate_fn=peptide_graph_encoder.masked_collate_fn)\n"
+    "    torch_dataset, batch_size=2, collate_fn=partial_collate_fn)\n"
    ]
   },
   {