yapf

mosaicml · Sep 16, 2024 · 103db99 · 103db99
1 parent 1977c0b
commit 103db99
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/llmfoundry/command_utils/data_prep/convert_text_to_mds.py b/llmfoundry/command_utils/data_prep/convert_text_to_mds.py
@@ -89,11 +89,10 @@ def __iter__(self) -> Iterable[dict[str, NDArray]]:
                     buffer += iids
                     while len(buffer) >= self.max_length:
                         concat_sample = buffer[:self.max_length]
-                        buffer = buffer[self.max_length:
-                                        ] if self.should_wrap else []
+                        buffer = buffer[self.
+                                        max_length:] if self.should_wrap else []
                         yield {
-                            'tokens':
-                                np.asarray(concat_sample, dtype=np.int32),
+                            'tokens': np.asarray(concat_sample, dtype=np.int32),
                         }
 
                         first_chunk = False