Safe

torchtitan/config_manager.py

@@ -270,6 +270,11 @@ def __init__(self):
             default=True,
             action="store_true",
             help="Whether to apply loss parallel when sequence parallel is enabled",
+        )        
+        self.parser.add_argument(
+            "--training.representation_type",
+            default="SMILES",
+            help="The representation type of the molecule for training the model.",
         )
         self.parser.add_argument(
             "--experimental.enable_async_tensor_parallel",

torchtitan/datasets/hf_datasets.py

@@ -87,6 +87,7 @@ def __init__(
         dataset_path: Optional[str],
         data_processing_style: str,
         tokenizer: Tokenizer,
+        representation_type: str = "SMILES",
         seq_len: int = 2048,
         world_size: int = 1,
         rank: int = 0,
@@ -133,6 +134,7 @@ def __init__(
         self._tokenizer = tokenizer
         self.seq_len = seq_len
         self.infinite = infinite
+        self.representation_type = representation_type
         self.rank = rank
         self.world_size = world_size
 
@@ -142,7 +144,6 @@ def __init__(
         else:
             self.store = None
 
-
         # variables for checkpointing
         self._sample_idx = 0
         self._all_tokens: List[int] = []
@@ -255,6 +256,7 @@ def build_hf_data_loader(
     seq_len: int,
     world_size,
     rank,
+    representation_type,
     infinite: bool = True,
     pin_memory: bool = False,
     num_workers: int = 2,

torchtitan/utils/dataset_utils.py

@@ -30,12 +30,12 @@ def load_jsonl_line(jsonl_line):
         raise ValueError(f"Error decoding JSON: {e}")
 
 
-def chemlactica_style_data_processing(sample_json, rng):
+def chemlactica_style_data_processing(sample_json, rng, representation_type):
     try:
         sample_json = json.loads(sample_json["text"])
         compound = delete_empty_tags(sample_json)
         sample_json = generate_formatted_string(
-            compound, rng
+            compound, rng, representation_type
         )
     except Exception as e:
         print(e)