just because omega starts with OMMMM does not mean it's zen

mosaicml · Apr 23, 2024 · 80acfb3 · 80acfb3
1 parent 08814e1
commit 80acfb3
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/llmfoundry/data/text_data.py b/llmfoundry/data/text_data.py
@@ -14,7 +14,6 @@
 import transformers
 from composer.core.data_spec import DataSpec
 from composer.core.types import Batch
-from omegaconf import OmegaConf as om
 from streaming import Stream, StreamingDataset
 from torch.utils.data import DataLoader
 from transformers import PreTrainedTokenizerBase
@@ -274,6 +273,7 @@ def build_text_dataloader(
     persistent_workers: bool = True,
     timeout: int = 0,
 ) -> DataSpec:
+
     dataset_cfg = dataset
 
     # get kwargs
@@ -450,7 +450,6 @@ def get_num_samples_in_batch(batch: Batch) -> int:
         'drop_last': False,
         'num_workers': 4,
     }
-    cfg = om.create(cfg)
     device_batch_size = 2
 
     tokenizer_name = args.tokenizer

diff --git a/tests/data/test_dataloader.py b/tests/data/test_dataloader.py
@@ -1119,7 +1119,7 @@ def test_token_counting_func_dataloader_setting(
                                          device_batch_size=batch_size,
                                          **cfg)
     elif dataloader_type == 'text':
-        cfg = DictConfig({
+        cfg = {
             'name': 'text',
             'dataset': {
                 'local': 'dummy-path',
@@ -1130,7 +1130,7 @@ def test_token_counting_func_dataloader_setting(
                 'shuffle_seed': 0,
             },
             **common_args
-        })
+        }
         ds_mock = MagicMock()
         ds_mock.tokenizer = gptt
         monkeypatch.setattr('llmfoundry.data.text_data.StreamingTextDataset',
@@ -1142,8 +1142,6 @@ def test_token_counting_func_dataloader_setting(
     else:
         raise NotImplementedError()
 
-    cfg = om.create(cfg)
-
     batch_collated = dl.dataloader.collate_fn(batch_tokenized)  # type: ignore
     actual_token_count = dl.get_num_tokens_in_batch(batch_collated)