unify spelling of Collator

Author: sfluegel

chebai/preprocessing/collate.py

@@ -5,7 +5,7 @@
 from chebai.preprocessing.structures import XYData
 
 
-class Collater:
+class Collator:
     """Base class for collating data samples into a batch."""
 
     def __init__(self, **kwargs):
@@ -23,8 +23,8 @@ def __call__(self, data: List[Dict]) -> XYData:
         raise NotImplementedError
 
 
-class DefaultCollater(Collater):
-    """Default collater that extracts features and labels."""
+class DefaultCollator(Collator):
+    """Default collator that extracts features and labels."""
 
     def __call__(self, data: List[Dict]) -> XYData:
         """Collate data samples by extracting features and labels.
@@ -39,11 +39,12 @@ def __call__(self, data: List[Dict]) -> XYData:
         return XYData(x, y)
 
 
-class RaggedCollater(Collater):
-    """Collater for handling ragged data samples."""
+class RaggedCollator(Collator):
+    """Collator for handling ragged data samples."""
 
     def __call__(self, data: List[Union[Dict, Tuple]]) -> XYData:
-        """Collate ragged data samples into a batch.
+        """Collate ragged data samples (i.e., samples of unequal size such as string representations of molecules) into
+        a batch.
 
         Args:
             data (List[Union[Dict, Tuple]]): List of ragged data samples.

chebai/preprocessing/datasets/pubchem.py

@@ -901,7 +901,7 @@ def dataloader(self, kind: str, **kwargs) -> DataLoader:
             unlabeled_data = unlabeled_data[: self.data_limit]
         return DataLoader(
             labeled_data + unlabeled_data,
-            collate_fn=self.reader.collater,
+            collate_fn=self.reader.collator,
             batch_size=self.batch_size,
             **kwargs,
         )

chebai/preprocessing/reader.py

@@ -1,12 +1,12 @@
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional
 
 from pysmiles.read_smiles import _tokenize
 from transformers import RobertaTokenizerFast
 import deepsmiles
 import selfies as sf
 
-from chebai.preprocessing.collate import DefaultCollater, RaggedCollater
+from chebai.preprocessing.collate import DefaultCollator, RaggedCollator
 
 EMBEDDING_OFFSET = 10
 PADDING_TOKEN_INDEX = 0
@@ -16,15 +16,16 @@
 
 class DataReader:
     """
-    Base class for reading and preprocessing data.
+    Base class for reading and preprocessing data. Turns the raw input data (e.g., a SMILES string) into the model
+    input format (e.g., a list of tokens).
 
     Args:
         collator_kwargs: Optional dictionary of keyword arguments for the collator.
         token_path: Optional path for the token file.
-        kwargs: Additional keyword arguments.
+        kwargs: Additional keyword arguments (not used).
     """
 
-    COLLATER = DefaultCollater
+    COLLATOR = DefaultCollator
 
     def __init__(
         self,
@@ -34,7 +35,7 @@ def __init__(
     ):
         if collator_kwargs is None:
             collator_kwargs = dict()
-        self.collater = self.COLLATER(**collator_kwargs)
+        self.collator = self.COLLATOR(**collator_kwargs)
         self.dirname = os.path.dirname(__file__)
         self._token_path = token_path
 
@@ -126,7 +127,7 @@ class ChemDataReader(DataReader):
         kwargs: Additional keyword arguments.
     """
 
-    COLLATER = RaggedCollater
+    COLLATOR = RaggedCollator
 
     @classmethod
     def name(cls) -> str:
@@ -201,7 +202,7 @@ class ChemDataUnlabeledReader(ChemDataReader):
         kwargs: Additional keyword arguments.
     """
 
-    COLLATER = RaggedCollater
+    COLLATOR = RaggedCollator
 
     @classmethod
     def name(cls) -> str:
@@ -220,13 +221,13 @@ class ChemBPEReader(DataReader):
     Args:
         data_path: Path for the pretrained BPE tokenizer.
         max_len: Maximum length of the tokenized sequence.
-        vsize: Vocabulary size for the tokenizer.
+        vsize: Vocabulary size for the tokenizer (not used).
         collator_kwargs: Optional dictionary of keyword arguments for the collator.
         token_path: Optional path for the token file.
         kwargs: Additional keyword arguments.
     """
 
-    COLLATER = RaggedCollater
+    COLLATOR = RaggedCollator
 
     @classmethod
     def name(cls) -> str:
@@ -264,7 +265,7 @@ class SelfiesReader(ChemDataReader):
         kwargs: Additional keyword arguments.
     """
 
-    COLLATER = RaggedCollater
+    COLLATOR = RaggedCollator
 
     def __init__(
         self,
@@ -309,7 +310,7 @@ class OrdReader(DataReader):
         kwargs: Additional keyword arguments.
     """
 
-    COLLATER = RaggedCollater
+    COLLATOR = RaggedCollator
 
     @classmethod
     def name(cls) -> str:

chebai/result/base.py

@@ -45,7 +45,7 @@ def _process_row(self, row):
 
     def _generate_predictions(self, data_path, raw=False, **kwargs):
         self._model.eval()
-        collate = self._reader.COLLATER()
+        collate = self._reader.COLLATOR()
         if raw:
             data_tuples = [
                 (x["features"], x["ident"], self._reader.to_data(self._process_row(x)))

chebai/result/pretraining.py

@@ -31,7 +31,7 @@ def evaluate_model(logs_base_path, model_filename, data_module):
         )
     )
     assert isinstance(model, electra.ElectraPre)
-    collate = data_module.reader.COLLATER()
+    collate = data_module.reader.COLLATOR()
     test_file = "test.pt"
     data_path = os.path.join(data_module.processed_dir, test_file)
     data_list = torch.load(data_path)

chebai/result/utils.py

@@ -73,7 +73,7 @@ def evaluate_model(
         Tensors with predictions and labels.
     """
     model.eval()
-    collate = data_module.reader.COLLATER()
+    collate = data_module.reader.COLLATOR()
 
     data_list = data_module.load_processed_data("test", filename)
     data_list = data_list[: data_module.data_limit]