Add OpReplaceAnyElements to ops_common

fuse/data/ops/ops_common.py

@@ -702,3 +702,61 @@ def __call__(
             raise Exception(f"Unsupported object type {type(input_obj)}")
 
         return sample_dict
+
+
+class OpReplaceAnyElements(OpBase):
+    """
+    Replace any element from a given list of elements with a requested value
+
+    For example, given the tensor [0,1,2,3,4,3,2,1,0], the list [0,2,4] and new value -100 - will return [-100,1,-100,3,-100,3,-100,1,-100].
+    This is useful, for example, when converting a set of token IDs to -100 to make a metric ignore certain elements (e.g. all special tokens).
+    """
+
+    def __call__(
+        self,
+        sample_dict: NDict,
+        key_in: str,
+        find_any_val: List[Any],
+        replace_with_val: Any,
+        key_out: str = None,
+    ) -> NDict:
+        """
+        Args:
+        key_in: the input numpy array, tensor, list or str
+        find_any_val: a list of values to will be replaces
+        replace_with_val: the value that will be used to replace "find_any_val" elements
+        key_out: if None, the modification will be inplace in "key_in" (faster), other wise, the name of the key to write to
+        """
+        assert (
+            key_in in sample_dict
+        ), f"Error: missing {key_in}, available keys {sample_dict.keypaths()} "
+
+        input_obj = sample_dict[key_in]
+
+        if key_out is not None:
+            if torch.is_tensor(input_obj):
+                input_obj = input_obj.clone()
+            elif isinstance(input_obj, np.ndarray):
+                input_obj = input_obj.copy()
+        else:
+            key_out = key_in
+
+        find_any_set = set(find_any_val)
+        if torch.is_tensor(input_obj) or isinstance(input_obj, np.ndarray):
+            output_obj = input_obj
+            output_obj[
+                (output_obj[..., None] == torch.tensor(find_any_val)).any(-1).nonzero()
+            ] = replace_with_val
+            sample_dict[key_out] = output_obj
+        elif isinstance(sample_dict[key_in], str):
+            sample_dict[key_out] = input_obj.translate(
+                str.maketrans({x: replace_with_val for x in find_any_val})
+            )
+        elif isinstance(input_obj, (list, tuple)):
+            sample_dict[key_out] = [
+                replace_with_val if x in find_any_set else x for x in input_obj
+            ]
+        else:
+            raise Exception(f"Unsupported object type {type(input_obj)}")
+
+        return sample_dict

fuse/data/ops/tests/test_ops_common.py

@@ -1,13 +1,22 @@
 import unittest
 
-from typing import Optional, OrderedDict, Union, List
+from typing import Optional, OrderedDict, Union, List, Any
+
+import torch
 
 from fuse.utils.ndict import NDict
 
 from fuse.data.ops.op_base import OpBase, OpReversibleBase, op_call
 from fuse.data.key_types_for_testing import DataTypeForTesting
 
-from fuse.data.ops.ops_common import OpApplyPatterns, OpFunc, OpLambda, OpRepeat
+from fuse.data.ops.ops_common import (
+    OpApplyPatterns,
+    OpFunc,
+    OpLambda,
+    OpRepeat,
+    OpReplaceAnyElements,
+    OpReplaceElements,
+)
 from fuse.data.ops.ops_common_for_testing import OpApplyTypesImaging
 
 
@@ -267,6 +276,88 @@ def test_op_apply_types(self) -> None:
         self.assertEqual(sample_dict["data.val.img_for_testing"], 4)
         self.assertEqual(sample_dict["model.a_seg_for_testing"], 2)
 
+    def test_op_replace_element(self) -> None:
+        """
+        test op_replace_element on each of possible inputs it supports: string, list, tensor
+        """
+
+        def _apply_op_replace(input: Any, to_replace: Any, replace_with: str) -> Any:
+            sample_dict = NDict({})
+            sample_dict["data.input"] = input
+            sample_dict = op_call(
+                OpReplaceElements(),
+                sample_dict,
+                "_.test_replace_element",
+                key_in="data.input",
+                key_out="data.out",
+                find_val=to_replace,
+                replace_with_val=replace_with,
+            )
+            return sample_dict["data.out"]
+
+        # input is a string
+        input = "abccba"
+        to_replace = "b"
+        replace_with = "X"
+        expected_result = "aXccXa"
+
+        res = _apply_op_replace(input, to_replace, replace_with)
+        self.assertEqual(res, expected_result)
+
+        # input is a list
+        input = [1, 2, 3, 2, 1]
+        to_replace = 2
+        replace_with = 0
+        expected_result = [1, 0, 3, 0, 1]
+        res = _apply_op_replace(input, to_replace, replace_with)
+        self.assertEqual(res, expected_result)
+
+        # input is a tensor
+        res = _apply_op_replace(torch.tensor(input), to_replace, replace_with)
+        self.assertTrue(torch.equal(res, torch.tensor(expected_result)))
+
+    def test_op_replace_any_element(self) -> None:
+        """
+        test op_replace_any_element on each of possible inputs it supports: string, list, tensor
+        """
+
+        def _apply_op_replace_any(
+            input: Any, to_replace: Any, replace_with: str
+        ) -> Any:
+            sample_dict = NDict({})
+            sample_dict["data.input"] = input
+            sample_dict = op_call(
+                OpReplaceAnyElements(),
+                sample_dict,
+                "_.test_replace_element",
+                key_in="data.input",
+                key_out="data.out",
+                find_any_val=to_replace,
+                replace_with_val=replace_with,
+            )
+            return sample_dict["data.out"]
+
+        # input is a string
+        input = "abccba"
+        to_replace = "bc"
+        replace_with = "X"
+        expected_result = "aXXXXa"
+
+        res = _apply_op_replace_any(input, to_replace, replace_with)
+        self.assertEqual(res, expected_result)
+
+        # input is a list
+        input = [1, 2, 3, 2, 1]
+        to_replace = [2, 3]
+        replace_with = 0
+        expected_result = [1, 0, 0, 0, 1]
+        res = _apply_op_replace_any(input, to_replace, replace_with)
+        self.assertEqual(res, expected_result)
+
+        # input is a tensor
+        res = _apply_op_replace_any(torch.tensor(input), to_replace, replace_with)
+        self.assertTrue(torch.equal(res, torch.tensor(expected_result)))
+
 
 if __name__ == "__main__":
     unittest.main()