Stream wrapper linter and code formatting

Signed-off-by: Alex-Brooks <Alex.Brooks@ibm.com>

Author: alex-jw-brooks

caikit_nlp/toolkit/data_stream_wrapper.py

@@ -16,14 +16,15 @@
 and objects for training / evaluating PyTorch models built around DataStreams, e.g., PyTorch
 DataLoaders, with minimal boilerplate.
 """
+# Standard
 from typing import Any, Iterator, Optional
 
 # Third Party
 from torch.utils.data import IterableDataset, get_worker_info
 
 # First Party
-from caikit.core.toolkit import error_handler
 from caikit.core.data_model import DataStream
+from caikit.core.toolkit import error_handler
 import alog
 
 log = alog.use_channel("STREAM_WRAP")
@@ -35,7 +36,13 @@ class SimpleIterableStreamWrapper(IterableDataset):
     compatability with PyTorch data loaders.
     """
 
-    def __init__(self, stream: DataStream[Any], shuffle: bool, buffer_size: Optional[int]=None, seed: int=42):
+    def __init__(
+        self,
+        stream: DataStream[Any],
+        shuffle: bool,
+        buffer_size: Optional[int] = None,
+        seed: int = 42,
+    ):
         error.type_check("<NLP12855513E>", bool, shuffle=shuffle)
         error.type_check(
             "<NLP12813713E>", int, buffer_size=buffer_size, allow_none=True
@@ -68,7 +75,7 @@ def __iter__(self) -> Iterator[Any]:
             # shuffles completed so far to ensure that every worker will
             # shuffle the same way for each epoch.
             shuffle_seed = self._get_shuffle_seed(worker_info)
-            log.debug(f"Reshuffling training data with seed: {shuffle_seed}")
+            log.debug("Reshuffling training data with seed: {}".format(shuffle_seed))
             cycle_stream = self.stream.shuffle(self.buffer_size, seed=shuffle_seed)
             self._increment_shuffle_seed(worker_info)
         else:
@@ -113,10 +120,9 @@ def _increment_shuffle_seed(self, worker_info: Optional["WorkerInfo"]) -> None:
         else:
             worker_info.dataset.shuffles_completed += 1
 
-    def _get_stream_partition(self,
-                              cycle_stream: DataStream[Any],
-                              worker_id: int,
-                              num_workers: int):
+    def _get_stream_partition(
+        self, cycle_stream: DataStream[Any], worker_id: int, num_workers: int
+    ):
         """Generator for a subset of a wrapped datastream; here, we simply traverse a stream,
         which is assumed to be preshuffled, and yield the elements that align with the
         scheme 'worker n gets every nth entry' after shuffling. This ensures that each
@@ -137,7 +143,7 @@ def _get_stream_partition(self,
 
     def __len__(self) -> int:
         """Gets the encapsulated stream length.
-        
+
         Returns:
             int
                 number of objects in the stream.

tests/toolkit/test_data_stream_wrapper.py

@@ -76,24 +76,25 @@ def test_iter_with_multi_worker():
     # Since we don't shuffle in this patched test, they should just be
     # divided as is.
     index_stream = [
-        {"label": 0}, # goes to worker 0
-        {"label": 1}, # goes to worker 1
-        {"label": 2}, # goes to worker 2
-        {"label": 3}, # goes to worker 0
-        {"label": 4}, # goes to worker 1
-        {"label": 5}, # goes to worker 2
+        {"label": 0},  # goes to worker 0
+        {"label": 1},  # goes to worker 1
+        {"label": 2},  # goes to worker 2
+        {"label": 3},  # goes to worker 0
+        {"label": 4},  # goes to worker 1
+        {"label": 5},  # goes to worker 2
     ]
     worker_info = [
         (w1_info, [index_stream[0], index_stream[3]]),
         (w2_info, [index_stream[1], index_stream[4]]),
         (w3_info, [index_stream[2], index_stream[5]]),
     ]
     for (dummy_worker, expected_elements) in worker_info:
-        with mock.patch.object(worker, '_worker_info', dummy_worker):
+        with mock.patch.object(worker, "_worker_info", dummy_worker):
             wrapper = SimpleIterableStreamWrapper(stream=index_stream, shuffle=False)
             for _ in range(NUM_CYCLES):
                 actual_elements = list(wrapper)
                 test_results.append(
-                    actual_elements == expected_elements and len(actual_elements) == len(expected_elements)
+                    actual_elements == expected_elements
+                    and len(actual_elements) == len(expected_elements)
                 )
     assert all(test_results)