Added experimental annotation to fixes #22564 (#22565)

apache · Aug 4, 2022 · c8d92b0 · c8d92b0
1 parent afb7dc9
commit c8d92b0
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/sdks/python/apache_beam/ml/inference/pytorch_inference.py b/sdks/python/apache_beam/ml/inference/pytorch_inference.py
@@ -29,6 +29,12 @@
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.ml.inference.base import ModelHandler
 from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.utils.annotations import experimental
+
+__all__ = [
+    'PytorchModelHandlerTensor',
+    'PytorchModelHandlerKeyedTensor',
+]
 
 
 def _load_model(
@@ -145,6 +151,7 @@ def get_metrics_namespace(self) -> str:
     return 'RunInferencePytorch'
 
 
+@experimental(extra_message="No backwards-compatibility guarantees.")
 class PytorchModelHandlerKeyedTensor(ModelHandler[Dict[str, torch.Tensor],
                                                   PredictionResult,
                                                   torch.nn.Module]):

diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference.py b/sdks/python/apache_beam/ml/inference/sklearn_inference.py
@@ -31,13 +31,19 @@
 from apache_beam.io.filesystems import FileSystems
 from apache_beam.ml.inference.base import ModelHandler
 from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.utils.annotations import experimental
 
 try:
   import joblib
 except ImportError:
   # joblib is an optional dependency.
   pass
 
+__all__ = [
+    'SklearnModelHandlerNumpy',
+    'SklearnModelHandlerPandas',
+]
+
 
 class ModelFileType(enum.Enum):
   """Defines how a model file is serialized. Options are pickle or joblib."""
@@ -132,6 +138,7 @@ def get_num_bytes(self, batch: Sequence[pandas.DataFrame]) -> int:
     return sum(sys.getsizeof(element) for element in batch)
 
 
+@experimental(extra_message="No backwards-compatibility guarantees.")
 class SklearnModelHandlerPandas(ModelHandler[pandas.DataFrame,
                                              PredictionResult,
                                              BaseEstimator]):