[data] update the arguments based on recent comment

Signed-off-by: Xingyu Long <xingyulong97@gmail.com>

Author: xingyu-long

python/ray/data/preprocessor.py

@@ -4,15 +4,7 @@
 import pickle
 import warnings
 from enum import Enum
-from typing import (
-    TYPE_CHECKING,
-    Any,
-    Dict,
-    Union,
-    List,
-    Optional,
-    Literal,
-)
+from typing import TYPE_CHECKING, Any, Dict, Union, List, Optional
 
 from ray.air.util.data_batch_conversion import BatchFormat
 from ray.util.annotations import DeveloperAPI, PublicAPI
@@ -132,7 +124,7 @@ def fit_transform(
         *,
         transform_num_cpus: Optional[float] = None,
         transform_memory: Optional[float] = None,
-        transform_batch_size: Union[int, None, Literal["default"]] = None,
+        transform_batch_size: Optional[int] = None,
         transform_concurrency: Optional[int] = None,
     ) -> "Dataset":
         """Fit this Preprocessor to the Dataset and then transform the Dataset.
@@ -143,10 +135,10 @@ def fit_transform(
 
         Args:
             ds: Input Dataset.
-            transform_num_cpus: The number of CPUs to reserve for each parallel map worker.
-            transform_memory: The heap memory in bytes to reserve for each parallel map worker.
-            transform_batch_size: The maximum number of rows to return.
-            transform_concurrency: The maximum number of Ray workers to use concurrently.
+            transform_num_cpus: [experimental] The number of CPUs to reserve for each parallel map worker.
+            transform_memory: [experimental] The heap memory in bytes to reserve for each parallel map worker.
+            transform_batch_size: [experimental] The maximum number of rows to return.
+            transform_concurrency: [experimental] The maximum number of Ray workers to use concurrently.
 
         Returns:
             ray.data.Dataset: The transformed Dataset.
@@ -164,19 +156,19 @@ def transform(
         self,
         ds: "Dataset",
         *,
+        batch_size: Optional[int] = None,
         num_cpus: Optional[float] = None,
         memory: Optional[float] = None,
-        batch_size: Union[int, None, Literal["default"]] = None,
         concurrency: Optional[int] = None,
     ) -> "Dataset":
         """Transform the given dataset.
 
         Args:
             ds: Input Dataset.
-            num_cpus: The number of CPUs to reserve for each parallel map worker.
-            memory: The heap memory in bytes to reserve for each parallel map worker.
-            batch_size: The maximum number of rows to return.
-            concurrency: The maximum number of Ray workers to use concurrently.
+            batch_size: [experimental] Advanced configuration for adjusting input size for each worker.
+            num_cpus: [experimental] The number of CPUs to reserve for each parallel map worker.
+            memory: [experimental] The heap memory in bytes to reserve for each parallel map worker.
+            concurrency: [experimental] The maximum number of Ray workers to use concurrently.
 
         Returns:
             ray.data.Dataset: The transformed Dataset.
@@ -193,7 +185,13 @@ def transform(
                 "`fit` must be called before `transform`, "
                 "or simply use fit_transform() to run both steps"
             )
-        transformed_ds = self._transform(ds, num_cpus, memory, batch_size, concurrency)
+        transformed_ds = self._transform(
+            ds,
+            batch_size=batch_size,
+            num_cpus=num_cpus,
+            memory=memory,
+            concurrency=concurrency,
+        )
         return transformed_ds
 
     def transform_batch(self, data: "DataBatchType") -> "DataBatchType":
@@ -258,9 +256,9 @@ def _determine_transform_to_use(self) -> BatchFormat:
     def _transform(
         self,
         ds: "Dataset",
+        batch_size: Optional[int],
         num_cpus: Optional[float] = None,
         memory: Optional[float] = None,
-        batch_size: Union[int, None, Literal["default"]] = None,
         concurrency: Optional[int] = None,
     ) -> "Dataset":
         transform_type = self._determine_transform_to_use()

python/ray/data/preprocessors/tokenizer.py

@@ -1,4 +1,4 @@
-from typing import Callable, List, Optional, Literal, Union
+from typing import Callable, List, Optional
 
 import pandas as pd