Blacking

Author: daviswer

fms_fsdp/config/training.py

@@ -15,7 +15,9 @@ class train_config:
     file_type: str = "arrow"
     col_name: str = "tokens"
     tokenizer_path: str = "/fsx/tokenizer"
-    datasets: str = "lang=en/dataset=commoncrawl,lang=en/dataset=webhose,lang=en/dataset=github_clean,lang=de/dataset=wikipedia,lang=es/dataset=wikipedia,lang=fr/dataset=wikipedia,lang=ja/dataset=wikipedia,lang=pt/dataset=wikipedia,lang=en/dataset=wikimedia,lang=en/dataset=uspto,lang=en/dataset=pubmedcentral,lang=en/dataset=arxiv,lang=en/dataset=stackexchange"
+    datasets: str = (
+        "lang=en/dataset=commoncrawl,lang=en/dataset=webhose,lang=en/dataset=github_clean,lang=de/dataset=wikipedia,lang=es/dataset=wikipedia,lang=fr/dataset=wikipedia,lang=ja/dataset=wikipedia,lang=pt/dataset=wikipedia,lang=en/dataset=wikimedia,lang=en/dataset=uspto,lang=en/dataset=pubmedcentral,lang=en/dataset=arxiv,lang=en/dataset=stackexchange"
+    )
     weights: str = "7725,500,550,28,17,22,25,8,100,500,175,250,100"
     seq_length: int = 4096
     vocab_size: int = 32000

fms_fsdp/utils/dataloader_utils.py

@@ -177,4 +177,4 @@ def splitstrip(x):
     datas = splitstrip(datas)
     weights = [float(x) for x in splitstrip(weights)]
     cols = splitstrip(cols)
-    return datas, weights, cols
+    return datas, weights, cols

fms_fsdp/utils/dataset_utils.py

@@ -355,14 +355,18 @@ def length(self, path: str):
         return self.open(path).num_record_batches
 
     def get(self, reader: pa.RecordBatchFileReader, index: int, drop_tokens: Set):
-        assert index < reader.num_record_batches, f"Illegal index {index} in set of {reader.num_record_batches} documents"
+        assert (
+            index < reader.num_record_batches
+        ), f"Illegal index {index} in set of {reader.num_record_batches} documents"
         frame = reader.get_batch(index)
         doc = None
         for name in self.col_names:
             if name in frame.column_names:
                 doc = frame[name]
                 break
-        assert doc is not None, f"None of column names {self.col_names} found in file headers {frame.column_names}"
+        assert (
+            doc is not None
+        ), f"None of column names {self.col_names} found in file headers {frame.column_names}"
         if len(doc) > 0 and doc[0].as_py() in drop_tokens:
             doc = doc.slice(1, len(doc) - 1)
         # Recheck len for edge case where doc=[eos]
@@ -382,7 +386,9 @@ class ParquetHandler(_ShardFileHandler):
     before getting/slicing. However, this is a standard and widely-used data format.
     """
 
-    def __init__(self, tokenizer_path: str, col_names: List[str] = ["text", "contents", "tokens"]):
+    def __init__(
+        self, tokenizer_path: str, col_names: List[str] = ["text", "contents", "tokens"]
+    ):
         self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
         self.col_names = col_names
 
@@ -396,14 +402,18 @@ def open(self, path: str):
             if name in names:
                 match = name
                 break
-        assert match is not None, f"None of column names {self.col_names} found in file headers {names}"
+        assert (
+            match is not None
+        ), f"None of column names {self.col_names} found in file headers {names}"
         return pq.read_pandas(path, columns=[match], partitioning=None)[match]
 
     def length(self, path: str):
         return pq.read_metadata(path).num_rows
 
     def get(self, reader, index: int, drop_tokens: Set):
-        assert index < reader.length(), f"Illegal index {index} in set of {reader.length()} documents"
+        assert (
+            index < reader.length()
+        ), f"Illegal index {index} in set of {reader.length()} documents"
         doc = self.tokenizer(str(reader[index])[:1_000_000])["input_ids"]
         if len(doc) > 0 and doc[0] in drop_tokens:
             doc = doc[1:]
@@ -417,7 +427,9 @@ def slice(self, doc: List, index: int, n_pull: int) -> List:
 
 
 class AutoHandler(_ShardFileHandler):
-    def __init__(self, tokenizer_path: str, col_names: List[str] = ["text", "contents", "tokens"]):
+    def __init__(
+        self, tokenizer_path: str, col_names: List[str] = ["text", "contents", "tokens"]
+    ):
         self.PHandler = ParquetHandler(tokenizer_path, col_names)
         self.AHandler = ArrowHandler(col_names)
         self.current = _ShardFileHandler()
@@ -1132,7 +1144,9 @@ def setup(self):
                 ndocs = doc_counts[shard]
                 if ndocs > 0:
                     doc_start = int(ndocs * shardset[shard][0])
-                    doc_end = max(doc_start, int(ndocs * shardset[shard][1]) - 1)  # inclusive upper bound
+                    doc_end = max(
+                        doc_start, int(ndocs * shardset[shard][1]) - 1
+                    )  # inclusive upper bound
                     self.docset.append([shard, doc_start, doc_end])
                     doccount += doc_end - doc_start + 1
             self._len = doccount
@@ -1280,7 +1294,9 @@ def __iter__(self):
                     yield self._construct_chunk(j, doc, n_chunks)
 
             # Check that epoch was non-empty
-            assert self.has_yielded, f"Empty logical shard detected: {self.dataset, self.docset}"
+            assert (
+                self.has_yielded
+            ), f"Empty logical shard detected: {self.dataset, self.docset}"
 
     def load_state_dict(self, state_dicts, sharded_input=False):
         self.setup()
@@ -1378,7 +1394,7 @@ def setup(self):
             assert (
                 sum(self.n_docs_remaining) > 0
             ), f"No documents detected in shard {self.rank} of {self.datapath}"
-                
+
             self.generator = torch.Generator().manual_seed(self.rank)
 
     def __iter__(self):
@@ -1487,7 +1503,9 @@ def __init__(
         )
         assert len(self.datasets) > 0, "You must specify at least one dataset"
         for d in datasets:
-            assert os.path.exists(os.path.join(datapath, d)), f"Invalid subdataset path: {os.path.join(datapath, d)}"
+            assert os.path.exists(
+                os.path.join(datapath, d)
+            ), f"Invalid subdataset path: {os.path.join(datapath, d)}"
 
         if weights is not None:
             assert len(weights) == len(