Upload split_utils.py with huggingface_hub

split_utils.py

@@ -1,8 +1,10 @@
 import itertools
-import random
 import re
+from typing import Dict
 
 from .generator_utils import ReusableGenerator
+from .random_utils import nested_seed
+from .stream import Stream
 
 
 def parse_random_mix_string(input_str):
@@ -67,7 +69,7 @@ def parse_slices_string(input_str):
     result_dict = {}
 
     # Split the input string into a list of sources
-    sources = re.split("\+", input_str)
+    sources = re.split(r"\+", input_str)
     for source in sources:
         # If the source has a slice, parse it
         match = re.fullmatch(r"(\w+)\[(\d*):(\d*)\]", source)
@@ -119,7 +121,7 @@ def slice_streams(input_streams, mapping):
               the new streams, which consist of parts of the old streams chained together.
 
     Raises:
-        ValueError: If a stream is supposed to be sliced at an index greater than its length.
+        ValueError: If a stream is supposed to be sliced at an index greater than its length or a negative one.
 
     Example:
         >>> old_streams = {"train": [1, 2, 3, 4, 5, 6, 7, 8, 9], "test": [10, 11, 12, 13, 14]}
@@ -205,15 +207,30 @@ def build_stream_routing(mapping):
     return stream_mapping
 
 
-def random_mix_generator(new_stream_name, new_stream_sources, stream_routing, rand, input_streams):
+def rename_split(input_streams: Dict[str, Stream], mapping: Dict[str, str]):
+    """
+    Renames the streams
+    Args:
+        input_streams (dict): A dictionary containing the input streams, where each key is
+                              the name of the stream and the value is an iterable or generator
+                              representing the stream.
+
+        mapping (dict): A dictionary specifying the mapping of old streams to new streams.
+
+    Returns:
+        dict: A dictionary containing the generated new streams, where each key is the name
+              of the new stream and the value is a generator representing the stream."""
+    return {mapping.get(key, key): val for key, val in input_streams.items()}
+
+
+def random_mix_generator(new_stream_name, new_stream_sources, stream_routing, input_streams):
     for old_stream_name in new_stream_sources:
         optinal_streams, weights = stream_routing[old_stream_name]
-        rand.seed(old_stream_name)
-
-        for item in input_streams[old_stream_name]:
-            choice = rand.choices(optinal_streams, weights=weights, k=1)[0]
-            if choice == new_stream_name:
-                yield item
+        with nested_seed(old_stream_name) as rand:
+            for item in input_streams[old_stream_name]:
+                choice = rand.choices(optinal_streams, weights=weights, k=1)[0]
+                if choice == new_stream_name:
+                    yield item
 
 
 def random_mix_streams(input_streams, mapping):
@@ -263,20 +280,18 @@ def random_mix_streams(input_streams, mapping):
     # Build stream routing
     stream_routing = build_stream_routing(mapping)
 
-    rand = random.Random()
-
-    # Create new stream generators
-    for new_stream_name, new_stream_sources in mapping.items():
-        new_streams[new_stream_name] = ReusableGenerator(
-            random_mix_generator,
-            gen_kwargs={
-                "new_stream_name": new_stream_name,
-                "new_stream_sources": new_stream_sources,
-                "stream_routing": stream_routing,
-                "rand": rand,
-                "input_streams": input_streams,
-            },
-        )
+    with nested_seed():
+        # Create new stream generators
+        for new_stream_name, new_stream_sources in mapping.items():
+            new_streams[new_stream_name] = ReusableGenerator(
+                random_mix_generator,
+                gen_kwargs={
+                    "new_stream_name": new_stream_name,
+                    "new_stream_sources": new_stream_sources,
+                    "stream_routing": stream_routing,
+                    "input_streams": input_streams,
+                },
+            )
 
     return new_streams