revert cast, lint

Author: abhi-mosaic

composer/datasets/__init__.py

@@ -21,13 +21,13 @@
 """
 from composer.datasets.ade20k import ADE20kDatasetHparams as ADE20kDatasetHparams
 from composer.datasets.brats import BratsDatasetHparams as BratsDatasetHparams
+from composer.datasets.c4 import C4DatasetHparams as C4DatasetHparams
 from composer.datasets.cifar10 import CIFAR10DatasetHparams as CIFAR10DatasetHparams
 from composer.datasets.dataloader import DataloaderHparams as DataloaderHparams
 from composer.datasets.dataloader import WrappedDataLoader as WrappedDataLoader
 from composer.datasets.dataset_registry import get_dataset_registry as get_dataset_registry
 from composer.datasets.evaluator import EvaluatorHparams as EvaluatorHparams
 from composer.datasets.glue import GLUEHparams as GLUEHparams
-from composer.datasets.c4 import C4DatasetHparams as C4DatasetHparams
 from composer.datasets.hparams import DatasetHparams as DatasetHparams
 from composer.datasets.hparams import SyntheticHparamsMixin as SyntheticHparamsMixin
 from composer.datasets.imagenet import ImagenetDatasetHparams as ImagenetDatasetHparams

composer/datasets/c4.py

@@ -4,19 +4,35 @@
 from dataclasses import dataclass
 from functools import partial
 from itertools import chain
-from typing import List, Optional
+from typing import List
 
 import yahp as hp
 from torch.utils.data import IterableDataset, get_worker_info
 
 from composer.core.types import Batch, DataSpec
 from composer.datasets.dataloader import DataloaderHparams
-from composer.datasets.hparams import DatasetHparams
 from composer.utils import dist
 
 log = logging.getLogger(__name__)
 
 
+def _split_dict_fn(batch: Batch, n_microbatches: int) -> List[Batch]:
+    if isinstance(batch, dict):
+        chunked = {k: v.chunk(n_microbatches) for k, v in batch.items()}
+        for k, v in chunked.items():
+            if len(v) != n_microbatches:
+                raise ValueError(
+                    f"Unable to split batch into microbatches. "
+                    f"Key '{k}' has chunked list: {v} with length {len(v)}, but expected length {n_microbatches}. ")
+        microbatches = []
+        for idx in range(n_microbatches):
+            mb = {k: v[idx] for k, v in chunked.items()}
+            microbatches.append(mb)
+        return microbatches
+    else:
+        raise ValueError(f'Expected batch to be of type Dict[str, Tensor], but got {type(batch)}')
+
+
 @dataclass
 class C4DatasetHparams(hp.Hparams):
     """Builds a DataSpec for the C4 (Colossal Cleaned CommonCrawl) dataset.
@@ -64,23 +80,13 @@ def validate(self):
         if self.mlm and self.mlm_probability <= 0:
             raise ValueError("Must provide a positive 'mlm_probability' when using masked language modeling.")
 
-    def _split_dict_fn(batch: Batch, n_microbatches: int) -> List[Batch]:
-        if isinstance(batch, dict):
-            chunked = {k: v.chunk(n_microbatches) for k, v in batch.items()}
-            for k, v in chunked.items():
-                if len(v) != n_microbatches:
-                    raise ValueError(
-                        f"Unable to split batch into microbatches. "
-                        f"Key '{k}' has chunked list: {v} with length {len(v)}, but expected length {n_microbatches}. ")
-            microbatches = []
-            for idx in range(n_microbatches):
-                mb = {k: v[idx] for k, v in chunked.items()}
-                microbatches.append(mb)
-            return microbatches
-        else:
-            raise ValueError(f'Expected batch to be of type Dict[str, Tensor], but got {type(batch)}')
-
     def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHparams) -> DataSpec:
+        try:
+            import transformers
+        except ImportError:
+            raise ImportError('HuggingFace transformers not installed. '
+                              'Please install with `pip install composer[nlp]`')
+
         # Get C4 dataset
         c4_dataset = C4Dataset(split=self.split,
                                max_samples=self.max_samples,
@@ -103,7 +109,7 @@ def initialize_object(self, batch_size: int, dataloader_hparams: DataloaderHpara
             drop_last=self.drop_last,
             collate_fn=collate_fn,
         ),
-                        split_batch=self._split_dict_fn)
+                        split_batch=_split_dict_fn)
 
 
 class C4Dataset(IterableDataset):

composer/datasets/dataset_registry.py

@@ -2,9 +2,9 @@
 
 from composer.datasets.ade20k import ADE20kDatasetHparams
 from composer.datasets.brats import BratsDatasetHparams
+from composer.datasets.c4 import C4DatasetHparams
 from composer.datasets.cifar10 import CIFAR10DatasetHparams
 from composer.datasets.glue import GLUEHparams
-from composer.datasets.c4 import C4DatasetHparams
 from composer.datasets.imagenet import ImagenetDatasetHparams
 from composer.datasets.lm_datasets import LMDatasetHparams
 from composer.datasets.mnist import MNISTDatasetHparams

composer/trainer/trainer_hparams.py

@@ -8,7 +8,7 @@
 import textwrap
 import warnings
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Dict, List, Optional
+from typing import TYPE_CHECKING, Dict, List, Optional, cast
 
 import yahp as hp
 
@@ -306,8 +306,8 @@ def validate(self):
         super().validate()
 
         if self.deepspeed is not None:
-            self.deepspeed["zero_stage"] = self.deepspeed.get("zero_stage", 0)
-            self.deepspeed["steps_per_print"] = self.deepspeed.get("steps_per_print", 1e20)
+            self.deepspeed["zero_stage"] = cast(int, self.deepspeed.get("zero_stage", 0))
+            self.deepspeed["steps_per_print"] = cast(int, self.deepspeed.get("steps_per_print", 1e20))
 
             if self.deterministic_mode and self.deepspeed["zero_stage"] > 0:
                 raise ValueError("Deepspeed with zero stage > 0 is not compatible with deterministic mode")