Merge pull request #27 from YerevaNN/revert-26-safe

Revert "Safe"
YerevaNN · Sep 26, 2024 · d584b5a · d584b5a
2 parents 76bdcf4 + e608677
commit d584b5a
Show file tree

Hide file tree

Showing 7 changed files with 4 additions and 485 deletions.
diff --git a/torchtitan/config_manager.py b/torchtitan/config_manager.py
@@ -270,11 +270,6 @@ 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",

diff --git a/torchtitan/datasets/hf_datasets.py b/torchtitan/datasets/hf_datasets.py
@@ -87,7 +87,6 @@ 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,
@@ -134,7 +133,6 @@ 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
 
@@ -144,6 +142,7 @@ def __init__(
         else:
             self.store = None
 
+
         # variables for checkpointing
         self._sample_idx = 0
         self._all_tokens: List[int] = []
@@ -256,7 +255,6 @@ def build_hf_data_loader(
     seq_len: int,
     world_size,
     rank,
-    representation_type,
     infinite: bool = True,
     pin_memory: bool = False,
     num_workers: int = 2,

diff --git a/torchtitan/utils/dataset_utils.py b/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, representation_type):
+def chemlactica_style_data_processing(sample_json, rng):
     try:
         sample_json = json.loads(sample_json["text"])
         compound = delete_empty_tags(sample_json)
         sample_json = generate_formatted_string(
-            compound, rng, representation_type
+            compound, rng
         )
     except Exception as e:
         print(e)