Upload stream.py with huggingface_hub

stream.py

@@ -3,10 +3,11 @@ from typing import Dict, Iterable
 
 from datasets import Dataset, DatasetDict, IterableDataset, IterableDatasetDict
 
-from .generator_utils import ReusableGenerator
+from .dataclass import Dataclass, OptionalField
+from .generator_utils import CopyingReusableGenerator, ReusableGenerator
 
 
-class Stream:
+class Stream(Dataclass):
     """A class for handling streaming data in a customizable way.
 
     This class provides methods for generating, caching, and manipulating streaming data.
@@ -14,24 +15,13 @@ class Stream:
     Attributes:
         generator (function): A generator function for streaming data.
         gen_kwargs (dict, optional): A dictionary of keyword arguments for the generator function.
-        streaming (bool): Whether the data is streaming or not.
         caching (bool): Whether the data is cached or not.
     """
 
-    def __init__(self, generator, gen_kwargs=None, streaming=True, caching=False):
-        """Initializes the Stream with the provided parameters.
-
-        Args:
-            generator (function): A generator function for streaming data.
-            gen_kwargs (dict, optional): A dictionary of keyword arguments for the generator function. Defaults to None.
-            streaming (bool, optional): Whether the data is streaming or not. Defaults to True.
-            caching (bool, optional): Whether the data is cached or not. Defaults to False.
-        """
-
-        self.generator = generator
-        self.gen_kwargs = gen_kwargs if gen_kwargs is not None else {}
-        self.streaming = streaming
-        self.caching = caching
+    generator: callable
+    gen_kwargs: Dict[str, any] = OptionalField(default_factory=dict)
+    caching: bool = False
+    copying: bool = False
 
     def _get_initator(self):
         """Private method to get the correct initiator based on the streaming and caching attributes.
@@ -39,16 +29,13 @@ class Stream:
         Returns:
             function: The correct initiator function.
         """
-        if self.streaming:
-            if self.caching:
-                return IterableDataset.from_generator
-            else:
-                return ReusableGenerator
+        if self.caching:
+            return Dataset.from_generator
         else:
-            if self.caching:
-                return Dataset.from_generator
+            if self.copying:
+                return CopyingReusableGenerator
             else:
-                raise ValueError("Cannot create non-streaming non-caching stream")
+                return ReusableGenerator
 
     def _get_stream(self):
         """Private method to get the stream based on the initiator function.
@@ -58,18 +45,9 @@ class Stream:
         """
         return self._get_initator()(self.generator, gen_kwargs=self.gen_kwargs)
 
-    def set_caching(self, caching):
-        self.caching = caching
-
-    def set_streaming(self, streaming):
-        self.streaming = streaming
-
     def __iter__(self):
         return iter(self._get_stream())
 
-    def unwrap(self):
-        return self._get_stream()
-
     def peak(self):
         return next(iter(self))
 
@@ -79,17 +57,6 @@ class Stream:
                 break
             yield instance
 
-    def __repr__(self):
-        return f"{self.__class__.__name__}(generator={self.generator.__name__}, gen_kwargs={self.gen_kwargs}, streaming={self.streaming}, caching={self.caching})"
-
-
-def is_stream(obj):
-    return isinstance(obj, IterableDataset) or isinstance(obj, Stream) or isinstance(obj, Dataset)
-
-
-def iterable_starter(iterable):
-    return iter(deepcopy(iterable))
-
 
 class MultiStream(dict):
     """A class for handling multiple streams of data in a dictionary-like format.
@@ -125,8 +92,13 @@ class MultiStream(dict):
         """
         yield from self[key]
 
-    def unwrap(self, cls):
-        return cls({key: value.unwrap() for key, value in self.items()})
+    def set_caching(self, caching: bool):
+        for stream in self.values():
+            stream.caching = caching
+
+    def set_copying(self, copying: bool):
+        for stream in self.values():
+            stream.copying = copying
 
     def to_dataset(self) -> DatasetDict:
         return DatasetDict(
@@ -144,12 +116,11 @@ class MultiStream(dict):
         super().__setitem__(key, value)
 
     @classmethod
-    def from_generators(cls, generators: Dict[str, ReusableGenerator], streaming=True, caching=False):
+    def from_generators(cls, generators: Dict[str, ReusableGenerator], caching=False, copying=False):
         """Creates a MultiStream from a dictionary of ReusableGenerators.
 
         Args:
             generators (Dict[str, ReusableGenerator]): A dictionary of ReusableGenerators.
-            streaming (bool, optional): Whether the data should be streaming or not. Defaults to True.
             caching (bool, optional): Whether the data should be cached or not. Defaults to False.
 
         Returns:
@@ -160,22 +131,21 @@ class MultiStream(dict):
         return cls(
             {
                 key: Stream(
-                    generator.get_generator(),
-                    gen_kwargs=generator.get_gen_kwargs(),
-                    streaming=streaming,
+                    generator.generator,
+                    gen_kwargs=generator.gen_kwargs,
                     caching=caching,
+                    copying=copying,
                 )
                 for key, generator in generators.items()
             }
         )
 
     @classmethod
-    def from_iterables(cls, iterables: Dict[str, Iterable], streaming=True, caching=False):
+    def from_iterables(cls, iterables: Dict[str, Iterable], caching=False, copying=False):
         """Creates a MultiStream from a dictionary of iterables.
 
         Args:
             iterables (Dict[str, Iterable]): A dictionary of iterables.
-            streaming (bool, optional): Whether the data should be streaming or not. Defaults to True.
             caching (bool, optional): Whether the data should be cached or not. Defaults to False.
 
         Returns:
@@ -184,7 +154,11 @@ class MultiStream(dict):
 
         return cls(
             {
-                key: Stream(iterable_starter, gen_kwargs={"iterable": iterable}, streaming=streaming, caching=caching)
+                key: Stream(
+                    iterable.__iter__,
+                    caching=caching,
+                    copying=copying,
+                )
                 for key, iterable in iterables.items()
             }
         )