Upload splitters.py with huggingface_hub

splitters.py

@@ -1,3 +1,4 @@
+import itertools
 from dataclasses import field
 from typing import Dict, List, Optional
 
@@ -11,16 +12,24 @@ class Splitter(MultiStreamOperator):
     pass
 
 
-import random
-
+from .random_utils import random
 from .split_utils import (
     parse_random_mix_string,
     parse_slices_string,
     random_mix_streams,
+    rename_split,
     slice_streams,
 )
 
 
+class RenameSplits(Splitter):
+    mapper: Dict[str, str]
+
+    def process(self, multi_stream: MultiStream) -> MultiStream:
+        generators = rename_split(multi_stream, self.mapper)
+        return MultiStream(generators)
+
+
 class SplitRandomMix(Splitter):
     mix: Dict[str, str]
 
@@ -30,6 +39,35 @@ class SplitRandomMix(Splitter):
         return MultiStream.from_generators(generators, streaming=True)
 
 
+class SeparateSplit(Splitter):
+    """
+    Separates a split (e.g. train) into several splits (e.g. train1, train2)
+    sizes must indicate the size of every split except the last. If no size is give for the last split,
+     it includes all the examples not allocated to any split.
+    """
+
+    from_split: str
+    to_split_names: List[str]
+    to_split_sizes: List[int]
+
+    def verify(self):
+        assert (
+            len(self.to_split_names) == len(self.to_split_sizes)
+            or len(self.to_split_names) == len(self.to_split_sizes) + 1
+        ), f"Examples num should be specified to all or all but the last splits, instead given {len(self.to_split_names)} split names and {len(self.to_split_sizes)} split sizes. \n split names:{self.to_split_names} split sizes {self.to_split_sizes}"
+        return super().verify()
+
+    def process(self, multi_stream: MultiStream) -> MultiStream:
+        mapping = {key: {key: [(None, None)]} for key in multi_stream.keys() if key != self.from_split}
+        so_far = 0
+        for name, size in itertools.zip_longest(self.to_split_names, self.to_split_sizes):
+            mapping[name] = {self.from_split: [(so_far, size)]}
+            if size:
+                so_far += size
+        generators = slice_streams(multi_stream, mapping)
+        return MultiStream.from_generators(generators, streaming=True)
+
+
 class SliceSplit(Splitter):
     slices: Dict[str, str]