Upload operators.py with huggingface_hub

operators.py

@@ -1,22 +1,26 @@
+import uuid
+from abc import abstractmethod
+from copy import deepcopy
 from dataclasses import field
-from typing import Any, Dict, Generator, Iterable, List, Optional, Union
+from itertools import zip_longest
+from typing import Any, Dict, Generator, Iterable, List, Optional, Tuple, Union
 
-from .text_utils import nested_tuple_to_string
 from .artifact import Artifact, fetch_artifact
+from .dict_utils import dict_delete, dict_get, dict_set, is_subpath
 from .operator import (
     MultiStream,
     MultiStreamOperator,
+    PagedStreamOperator,
     SingleStreamOperator,
     SingleStreamReducer,
-    Stream,
+    StreamingOperator,
     StreamInitializerOperator,
     StreamInstanceOperator,
-    PagedStreamOperator,
 )
+from .random_utils import random
 from .stream import MultiStream, Stream
+from .text_utils import nested_tuple_to_string
 from .utils import flatten_dict
-import random
-from .utils import dict_query
 
 
 class FromIterables(StreamInitializerOperator):
@@ -26,20 +30,47 @@ class FromIterables(StreamInitializerOperator):
     Args:
         iterables (Dict[str, Iterable]): A dictionary where each key-value pair represents a stream name and its corresponding iterable.
     """
+
     def process(self, iterables: Dict[str, Iterable]) -> MultiStream:
         return MultiStream.from_iterables(iterables)
 
 
-class MapInstanceValues(StreamInstanceOperator):
+class RenameFields(StreamInstanceOperator):
+    """
+    Renames fields
+    Attributes:
+    mapper (Dict[str, str]): old field names to new field names dict
     """
-    Maps values in each instance of a stream based on the provided mappers.
 
-    Args:
-        mappers (Dict[str, Dict[str, str]]): A dictionary where each key-value pair represents a field in the instance and a mapper for that field.
-        strict (bool): If True, the operator will raise a KeyError if a value is not in its corresponding mapper. If False, unmapped values will be left unchanged. Defaults to True.
+    mapper: Dict[str, str]
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        result = {}
+        # passes on all values to preserve ordering
+        for key, value in instance.items():
+            result[self.mapper.get(key, key)] = value
+        # doesn't warn if unnecessary mapping was supplied for efficiency
+        return result
+
+
+class MapInstanceValues(StreamInstanceOperator):
+    """A class used to map instance values into a stream.
+
+    This class is a type of StreamInstanceOperator,
+    it maps values of instances in a stream using predefined mappers.
+
+    Attributes:
+        mappers (Dict[str, Dict[str, str]]): The mappers to use for mapping instance values.
+            Keys are the names of the fields to be mapped, and values are dictionaries
+            that define the mapping from old values to new values.
+        strict (bool): If True, the mapping is applied strictly. That means if a value
+            does not exist in the mapper, it will raise a KeyError. If False, values
+            that are not present in the mapper are kept as they are.
     """
+
     mappers: Dict[str, Dict[str, str]]
     strict: bool = True
+    use_query = False
 
     def verify(self):
         # make sure the mappers are valid
@@ -49,18 +80,16 @@ class MapInstanceValues(StreamInstanceOperator):
                 assert isinstance(k, str), f'Key "{k}" in mapper for field "{key}" should be a string, got {type(k)}'
 
     def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
-        result = {}
-        for key, value in instance.items():
-            str_value = str(value)
-            if key in self.mappers:
-                mapper = self.mappers[key]
+        for key, mapper in self.mappers.items():
+            value = dict_get(instance, key, use_dpath=self.use_query)
+            if value is not None:
+                value = str(value)  # make sure the value is a string
                 if self.strict:
-                    value = mapper[str_value]
+                    dict_set(instance, key, mapper[value], use_dpath=self.use_query)
                 else:
-                    if str_value in mapper:
-                        value = mapper[str_value]
-            result[key] = value
-        return result
+                    if value in mapper:
+                        dict_set(instance, key, mapper[value], use_dpath=self.use_query)
+        return instance
 
 
 class FlattenInstances(StreamInstanceOperator):
@@ -71,9 +100,10 @@ class FlattenInstances(StreamInstanceOperator):
         parent_key (str): A prefix to use for the flattened keys. Defaults to an empty string.
         sep (str): The separator to use when concatenating nested keys. Defaults to "_".
     """
+
     parent_key: str = ""
     sep: str = "_"
-    
+
     def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
         return flatten_dict(instance, parent_key=self.parent_key, sep=self.sep)
 
@@ -85,24 +115,244 @@ class AddFields(StreamInstanceOperator):
     Args:
         fields (Dict[str, object]): The fields to add to each instance.
     """
+
     fields: Dict[str, object]
+    use_query: bool = False
+    use_deepcopy: bool = False
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        if self.use_query:
+            for key, value in self.fields.items():
+                if self.use_deepcopy:
+                    value = deepcopy(value)
+                dict_set(instance, key, value, use_dpath=self.use_query)
+        else:
+            if self.use_deepcopy:
+                self.fields = deepcopy(self.fields)
+            instance.update(self.fields)
+        return instance
+
+
+class FieldOperator(StreamInstanceOperator):
+    """
+    A general stream that processes the values of a field (or multiple ones
+    Args:
+        field (Optional[str]): The field to process, if only a single one is passed Defaults to None
+        to_field (Optional[str]): Field name to save, if only one field is to be saved, if None is passed the operator would happen in-place and replace "field" Defaults to None
+        field_to_field (Optional[Union[List[Tuple[str, str]], Dict[str, str]]]): Mapping from fields to process to their names after this process, duplicates are allowed. Defaults to None
+        process_every_value (bool): Processes the values in a list instead of the list as a value, similar to *var. Defaults to False
+        use_query (bool): Whether to use dpath style queries. Defaults to False
+    """
+
+    field: Optional[str] = None
+    to_field: Optional[str] = None
+    field_to_field: Optional[Union[List[Tuple[str, str]], Dict[str, str]]] = None
+    process_every_value: bool = False
+    use_query: bool = False
+
+    def verify(self):
+        super().verify()
+
+        assert self.field is not None or self.field_to_field is not None, "Must supply a field to work on"
+        assert (
+            self.to_field is None or self.field_to_field is None
+        ), f"Can not apply operator to create both on {self.to_field} and on the mapping from fields to fields {self.field_to_field}"
+        assert (
+            self.field is None or self.field_to_field is None
+        ), f"Can not apply operator both on {self.field} and on the mapping from fields to fields {self.field_to_field}"
+        assert self._field_to_field, f"the from and to fields must be defined got: {self._field_to_field}"
+
+    @abstractmethod
+    def process_value(self, value: Any) -> Any:
+        pass
+
+    def prepare(self):
+        if self.to_field is None:
+            self.to_field = self.field
+        if self.field_to_field is None:
+            self._field_to_field = [(self.field, self.to_field)]
+        else:
+            try:
+                self._field_to_field = [(k, v) for k, v in self.field_to_field.items()]
+            except AttributeError:
+                self._field_to_field = self.field_to_field
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        for from_field, to_field in self._field_to_field:
+            old_value = dict_get(instance, from_field, use_dpath=self.use_query)
+            if self.process_every_value:
+                new_value = [self.process_value(value) for value in old_value]
+            else:
+                new_value = self.process_value(old_value)
+            if self.use_query and is_subpath(from_field, to_field):
+                dict_delete(instance, from_field)
+            dict_set(instance, to_field, new_value, use_dpath=self.use_query, not_exist_ok=True)
+        return instance
+
+
+class JoinStr(FieldOperator):
+    """
+    Joins a list of strings (contents of a field), similar to str.join()
+    Args:
+        separator (str): text to put between values
+    """
+
+    separator: str = ","
+
+    def process_value(self, value: Any) -> Any:
+        return self.separator.join(str(x) for x in value)
+
+
+class ZipFieldValues(StreamInstanceOperator):
+    """
+    Zips values of multiple fields similar to list(zip(*fields))
+    """
+
+    fields: str
+    to_field: str
+    longest: bool = False
+    use_query: bool = False
 
     def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
-        instance.update(self.fields)
+        values = []
+        for field in self.fields:
+            values.append(dict_get(instance, field, use_dpath=self.use_query))
+        if self.longest:
+            zipped = zip_longest(*values)
+        else:
+            zipped = zip(*values)
+        instance[self.to_field] = list(zipped)
         return instance
 
 
-class MapNestedDictValuesByQueries(StreamInstanceOperator):
-    field_to_query: Dict[str, str]
+class IndexOf(StreamInstanceOperator):
+    """
+    Finds the location of one value in another (iterable) value similar to to_field=search_in.index(index_of)
+    """
+
+    search_in: str
+    index_of: str
+    to_field: str
+    use_query: bool = False
 
     def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
-        updates = {}
-        for field, query in self.field_to_query.items():
-            updates[field] = dict_query(instance, query)
-        instance.update(updates)
+        lst = dict_get(instance, self.search_in, use_dpath=self.use_query)
+        item = dict_get(instance, self.index_of, use_dpath=self.use_query)
+        instance[self.to_field] = lst.index(item)
         return instance
 
 
+class TakeByField(StreamInstanceOperator):
+    """
+    Takes value from one field based on another field similar to field[index]
+    """
+
+    field: str
+    index: str
+    to_field: str = None
+    use_query: bool = False
+
+    def prepare(self):
+        if self.to_field is None:
+            self.to_field = self.field
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        value = dict_get(instance, self.field, use_dpath=self.use_query)
+        index_value = dict_get(instance, self.index, use_dpath=self.use_query)
+        instance[self.to_field] = value[index_value]
+        return instance
+
+
+class CopyFields(FieldOperator):
+    """
+    Copies specified fields from one field to another.
+
+    Args:
+        field_to_field (Union[List[List], Dict[str, str]]): A list of lists, where each sublist contains the source field and the destination field, or a dictionary mapping source fields to destination fields.
+        use_dpath (bool): Whether to use dpath for accessing fields. Defaults to False.
+    """
+
+    def process_value(self, value: Any) -> Any:
+        return value
+
+
+class AddID(StreamInstanceOperator):
+    id_field_name: str = "id"
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        instance[self.id_field_name] = str(uuid.uuid4()).replace("-", "")
+        return instance
+
+
+class CastFields(StreamInstanceOperator):
+    """
+    Casts specified fields to specified types.
+
+    Args:
+        types (Dict[str, str]): A dictionary mapping fields to their new types.
+        nested (bool): Whether to cast nested fields. Defaults to False.
+        fields (Dict[str, str]): A dictionary mapping fields to their new types.
+        defaults (Dict[str, object]): A dictionary mapping types to their default values for cases of casting failure.
+    """
+
+    types = {
+        "int": int,
+        "float": float,
+        "str": str,
+        "bool": bool,
+    }
+    fields: Dict[str, str] = field(default_factory=dict)
+    failure_defaults: Dict[str, object] = field(default_factory=dict)
+    use_nested_query: bool = False
+    cast_multiple: bool = False
+
+    def _cast_single(self, value, type, field):
+        try:
+            return self.types[type](value)
+        except:
+            if field not in self.failure_defaults:
+                raise ValueError(
+                    f'Failed to cast field "{field}" with value {value} to type "{type}", and no default value is provided.'
+                )
+            return self.failure_defaults[field]
+
+    def _cast_multiple(self, values, type, field):
+        values = [self._cast_single(value, type, field) for value in values]
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        for field, type in self.fields.items():
+            value = dict_get(instance, field, use_dpath=self.use_nested_query)
+            if self.cast_multiple:
+                casted_value = self._cast_multiple(value, type, field)
+            else:
+                casted_value = self._cast_single(value, type, field)
+            dict_set(instance, field, casted_value, use_dpath=self.use_nested_query)
+        return instance
+
+
+def recursive_divide(instance, divisor, strict=False):
+    if isinstance(instance, dict):
+        for key, value in instance.items():
+            instance[key] = recursive_divide(value, divisor, strict=strict)
+    elif isinstance(instance, list):
+        for i, value in enumerate(instance):
+            instance[i] = recursive_divide(value, divisor, strict=strict)
+    elif isinstance(instance, float):
+        instance /= divisor
+    elif strict:
+        raise ValueError(f"Cannot divide instance of type {type(instance)}")
+    return instance
+
+
+class DivideAllFieldsBy(StreamInstanceOperator):
+    divisor: float = 1.0
+    strict: bool = False
+    recursive: bool = True
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        return recursive_divide(instance, self.divisor, strict=self.strict)
+
+
 class ArtifactFetcherMixin:
     """
     Provides a way to fetch and cache artifacts in the system.
@@ -110,6 +360,7 @@ class ArtifactFetcherMixin:
     Args:
         cache (Dict[str, Artifact]): A cache for storing fetched artifacts.
     """
+
     cache: Dict[str, Artifact] = {}
 
     @classmethod
@@ -129,6 +380,7 @@ class ApplyValueOperatorsField(StreamInstanceOperator, ArtifactFetcherMixin):
         operators_field (str): The field containing the operators to be applied.
         default_operators (List[str]): A list of default operators to be used if no operators are found in the instance.
     """
+
     value_field: str
     operators_field: str
     default_operators: List[str] = None
@@ -158,6 +410,7 @@ class FilterByValues(SingleStreamOperator):
     Args:
         values (Dict[str, Any]): The values that instances should match to be included in the output.
     """
+
     values: Dict[str, Any]
 
     def process(self, stream: Stream, stream_name: str = None) -> Generator:
@@ -173,6 +426,7 @@ class Unique(SingleStreamReducer):
     Args:
         fields (List[str]): The fields that should be unique in each instance.
     """
+
     fields: List[str] = field(default_factory=list)
 
     @staticmethod
@@ -201,6 +455,7 @@ class SplitByValue(MultiStreamOperator):
     Args:
         fields (List[str]): The fields to use when splitting the MultiStream.
     """
+
     fields: List[str] = field(default_factory=list)
 
     def process(self, multi_stream: MultiStream) -> MultiStream:
@@ -227,6 +482,7 @@ class ApplyStreamOperatorsField(SingleStreamOperator, ArtifactFetcherMixin):
         field (str): The field containing the operators to be applied.
         reversed (bool): Whether to apply the operators in reverse order.
     """
+
     field: str
     reversed: bool = False
 
@@ -242,10 +498,9 @@ class ApplyStreamOperatorsField(SingleStreamOperator, ArtifactFetcherMixin):
 
         for operator_name in operators:
             operator = self.get_artifact(operator_name)
-            assert isinstance(
-                operator, SingleStreamOperator
-            ), f"Operator {operator_name} must be a SingleStreamOperator"
-            stream = operator.process(stream)
+            assert isinstance(operator, StreamingOperator), f"Operator {operator_name} must be a SingleStreamOperator"
+
+            stream = operator(MultiStream({"tmp": stream}))["tmp"]
 
         yield from stream
 
@@ -257,6 +512,7 @@ class AddFieldNamePrefix(StreamInstanceOperator):
     Args:
         prefix_dict (Dict[str, str]): A dictionary mapping stream names to prefixes.
     """
+
     prefix_dict: Dict[str, str]
 
     def prepare(self):
@@ -275,6 +531,7 @@ class MergeStreams(MultiStreamOperator):
         add_origin_stream_name (bool): Whether to add the origin stream name to each instance.
         origin_stream_name_field_name (str): The field name for the origin stream name.
     """
+
     new_stream_name: str = "all"
     add_origin_stream_name: bool = True
     origin_stream_name_field_name: str = "origin"
@@ -289,6 +546,7 @@ class MergeStreams(MultiStreamOperator):
     def process(self, multi_stream: MultiStream) -> MultiStream:
         return MultiStream({self.new_stream_name: Stream(self.merge, gen_kwargs={"multi_stream": multi_stream})})
 
+
 class Shuffle(PagedStreamOperator):
     """
     Shuffles the order of instances in each page of a stream.
@@ -296,6 +554,35 @@ class Shuffle(PagedStreamOperator):
     Args:
         page_size (int): The size of each page in the stream. Defaults to 1000.
     """
+
     def process(self, page: List[Dict], stream_name: str = None) -> Generator:
         random.shuffle(page)
-        yield from page
+        yield from page
+
+
+class EncodeLabels(StreamInstanceOperator):
+    """
+    Encode labels of specified fields together a into integers.
+
+    Args:
+        fields (List[str]): The fields to encode together.
+    """
+
+    fields: List[str]
+
+    def _process_multi_stream(self, multi_stream: MultiStream) -> MultiStream:
+        self.encoder = {}
+        return super()._process_multi_stream(multi_stream)
+
+    def process(self, instance: Dict[str, Any], stream_name: str = None) -> Dict[str, Any]:
+        for field in self.fields:
+            values = dict_get(instance, field, use_dpath=True)
+            if not isinstance(values, list):
+                values = [values]
+            for value in values:
+                if value not in self.encoder:
+                    self.encoder[value] = len(self.encoder)
+            new_values = [self.encoder[value] for value in values]
+            dict_set(instance, field, new_values, use_dpath=True, set_multiple=True)
+
+        return instance