chore: import upstream snapshot with attribution

python/ray/data/preprocessors/tokenizer.py

@@ -0,0 +1,142 @@
+from typing import Any, Callable, Dict, List, Optional
+
+import pandas as pd
+
+from ray.data.preprocessor import SerializablePreprocessorBase
+from ray.data.preprocessors.utils import (
+    _Computed,
+    _PublicField,
+    migrate_private_fields,
+    simple_split_tokenizer,
+)
+from ray.data.preprocessors.version_support import SerializablePreprocessor
+from ray.util.annotations import PublicAPI
+
+
+@PublicAPI(stability="alpha")
+@SerializablePreprocessor(version=1, identifier="io.ray.preprocessors.tokenizer")
+class Tokenizer(SerializablePreprocessorBase):
+    """Replace each string with a list of tokens.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import ray
+        >>> df = pd.DataFrame({"text": ["Hello, world!", "foo bar\\nbaz"]})
+        >>> ds = ray.data.from_pandas(df)  # doctest: +SKIP
+
+        The default ``tokenization_fn`` delimits strings using the space character.
+
+        >>> from ray.data.preprocessors import Tokenizer
+        >>> tokenizer = Tokenizer(columns=["text"])
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                       text
+        0  [Hello,, world!]
+        1   [foo, bar\\nbaz]
+
+        If the default logic isn't adequate for your use case, you can specify a
+        custom ``tokenization_fn``.
+
+        >>> import string
+        >>> def tokenization_fn(s):
+        ...     for character in string.punctuation:
+        ...         s = s.replace(character, "")
+        ...     return s.split()
+        >>> tokenizer = Tokenizer(columns=["text"], tokenization_fn=tokenization_fn)
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                      text
+        0   [Hello, world]
+        1  [foo, bar, baz]
+
+        :class:`Tokenizer` can also be used in append mode by providing the
+        name of the output_columns that should hold the tokenized values.
+
+        >>> tokenizer = Tokenizer(columns=["text"], output_columns=["text_tokenized"])
+        >>> tokenizer.transform(ds).to_pandas()  # doctest: +SKIP
+                    text    text_tokenized
+        0  Hello, world!  [Hello,, world!]
+        1   foo bar\\nbaz   [foo, bar\\nbaz]
+
+    Args:
+        columns: The columns to tokenize.
+        tokenization_fn: The function used to generate tokens. This function
+            should accept a string as input and return a list of tokens as
+            output. If unspecified, the tokenizer uses a function equivalent to
+            ``lambda s: s.split(" ")``.
+        output_columns: The names of the transformed columns. If None, the transformed
+            columns will be the same as the input columns. If not None, the length of
+            ``output_columns`` must match the length of ``columns``, othwerwise an error
+            will be raised.
+    """
+
+    _is_fittable = False
+
+    def __init__(
+        self,
+        columns: List[str],
+        tokenization_fn: Optional[Callable[[str], List[str]]] = None,
+        output_columns: Optional[List[str]] = None,
+    ):
+        super().__init__()
+        self._columns = columns
+        # TODO(matt): Add a more robust default tokenizer.
+        self._tokenization_fn = tokenization_fn or simple_split_tokenizer
+        self._output_columns = (
+            SerializablePreprocessorBase._derive_and_validate_output_columns(
+                columns, output_columns
+            )
+        )
+
+    @property
+    def columns(self) -> List[str]:
+        return self._columns
+
+    @property
+    def tokenization_fn(self) -> Callable[[str], List[str]]:
+        return self._tokenization_fn
+
+    @property
+    def output_columns(self) -> List[str]:
+        return self._output_columns
+
+    def _transform_pandas(self, df: pd.DataFrame):
+        def column_tokenizer(s: pd.Series):
+            return s.map(self._tokenization_fn)
+
+        df[self._output_columns] = df.loc[:, self._columns].transform(column_tokenizer)
+        return df
+
+    def __repr__(self):
+        name = getattr(self._tokenization_fn, "__name__", self._tokenization_fn)
+        return (
+            f"{self.__class__.__name__}(columns={self._columns!r}, "
+            f"tokenization_fn={name}, output_columns={self._output_columns!r})"
+        )
+
+    def _get_serializable_fields(self) -> Dict[str, Any]:
+        return {
+            "columns": self._columns,
+            "tokenization_fn": self._tokenization_fn,
+            "output_columns": self._output_columns,
+        }
+
+    def _set_serializable_fields(self, fields: Dict[str, Any], version: int):
+        # required fields
+        self._columns = fields["columns"]
+        self._tokenization_fn = fields["tokenization_fn"]
+        self._output_columns = fields["output_columns"]
+
+    def __setstate__(self, state: Dict[str, Any]) -> None:
+        super().__setstate__(state)
+        migrate_private_fields(
+            self,
+            fields={
+                "_columns": _PublicField(public_field="columns"),
+                "_tokenization_fn": _PublicField(
+                    public_field="tokenization_fn", default=simple_split_tokenizer
+                ),
+                "_output_columns": _PublicField(
+                    public_field="output_columns",
+                    default=_Computed(lambda obj: obj._columns),
+                ),
+            },
+        )