fmeval.transforms.batched_transform

 1import numpy as np
 2
 3from typing import Dict
 4from abc import abstractmethod
 5from fmeval.transforms.transform import Transform
 6
 7
 8class BatchedTransform(Transform):
 9    """A BatchedTransform is a Transform that takes in a batch of records instead of a single record.
10
11    Certain transforms will have a significant performance boost when processing records in batches
12    (the performance boost depends on the logic internal to the transform's __call__ method).
13
14    This abstract base class should be inherited by such transforms.
15    """
16
17    def __init__(self, *args, **kwargs):
18        super().__init__(*args, **kwargs)
19
20    @abstractmethod
21    def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
22        """Return a batch of records containing data that gets computed in this method.
23
24        :param batch: The batch to be transformed.
25        :returns: A batch of records containing data that gets computed in this method.
26            This batch can be the same object as the input batch. In this case,
27            the logic in this method should mutate the input batch directly.
28        """
29
30    @property
31    def batch_size(self) -> int:
32        """The size of the batches that this transform should process.
33
34        Defaults to -1, in which case default batch size options will
35        be used when executing the transform.
36        """
37        return -1  # pragma: no cover