Upload operator.py with huggingface_hub

operator.py

@@ -3,6 +3,7 @@ from dataclasses import field
 from typing import Any, Dict, Generator, List, Optional, Union
 
 from .artifact import Artifact
+from .dataclass import NonPositionalField
 from .random_utils import nested_seed
 from .stream import MultiStream, Stream
 
@@ -83,9 +84,14 @@ class SourceOperator(StreamSource):
 
     """
 
+    caching: bool = NonPositionalField(default=None)
+
     def __call__(self) -> MultiStream:
         with nested_seed():
-            return self.process()
+            multi_stream = self.process()
+            if self.caching is not None:
+                multi_stream.set_caching(self.caching)
+            return multi_stream
 
     @abstractmethod
     def process(self) -> MultiStream:
@@ -102,8 +108,13 @@ class StreamInitializerOperator(StreamSource):
 
     """
 
+    caching: bool = NonPositionalField(default=None)
+
     def __call__(self, *args, **kwargs) -> MultiStream:
         with nested_seed():
+            multi_stream = self.process(*args, **kwargs)
+            if self.caching is not None:
+                multi_stream.set_caching(self.caching)
             return self.process(*args, **kwargs)
 
     @abstractmethod
@@ -118,6 +129,8 @@ class MultiStreamOperator(StreamingOperator):
     A multi-stream operator is a type of `StreamingOperator` that operates on an entire MultiStream object at once. It takes a `MultiStream` as input and produces a `MultiStream` as output. The `process` method should be implemented by subclasses to define the specific operations to be performed on the input `MultiStream`.
     """
 
+    caching: bool = NonPositionalField(default=None)
+
     def __call__(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
         with nested_seed():
             return self._process_multi_stream(multi_stream)
@@ -125,6 +138,8 @@ class MultiStreamOperator(StreamingOperator):
     def _process_multi_stream(self, multi_stream: Optional[MultiStream] = None) -> MultiStream:
         result = self.process(multi_stream)
         assert isinstance(result, MultiStream), "MultiStreamOperator must return a MultiStream"
+        if self.caching is not None:
+            result.set_caching(self.caching)
         return result
 
     @abstractmethod
@@ -198,7 +213,7 @@ class SingleStreamReducer(StreamingOperator):
         return result
 
     @abstractmethod
-    def process(self, stream: Stream) -> Any:
+    def process(self, stream: Stream) -> Stream:
         pass
 
 
@@ -296,7 +311,7 @@ class InstanceOperatorWithGlobalAccess(StreamingOperator):
 
         if self.cache_accessible_streams:
             for stream in self.accessible_streams.values():
-                stream.set_caching(True)
+                stream.caching = True
 
         for stream_name, stream in multi_stream.items():
             stream = Stream(self.generator, gen_kwargs={"stream": stream, "multi_stream": self.accessible_streams})