Refactor base

sdks/python/apache_beam/examples/inference/pytorch_image_classification.py

@@ -27,8 +27,8 @@
 import apache_beam as beam
 import torch
 from apache_beam.io.filesystems import FileSystems
-from apache_beam.ml.inference.api import PredictionResult
-from apache_beam.ml.inference.api import RunInference
+from apache_beam.ml.inference.base import PredictionResult
+from apache_beam.ml.inference.base import RunInference
 from apache_beam.ml.inference.pytorch_inference import PytorchModelLoader
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.options.pipeline_options import SetupOptions

sdks/python/apache_beam/ml/inference/__init__.py

@@ -14,3 +14,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
+from apache_beam.ml.inference.base import RunInference

sdks/python/apache_beam/ml/inference/base.py

@@ -27,17 +27,21 @@
 Note: This module is still actively being developed and users should have
 no expectation that these interfaces will not change.
 """
+# mypy: ignore-errors
 
 import logging
 import pickle
 import sys
 import time
+from dataclasses import dataclass
 from typing import Any
 from typing import Generic
 from typing import Iterable
 from typing import List
 from typing import Mapping
+from typing import Tuple
 from typing import TypeVar
+from typing import Union
 
 import apache_beam as beam
 from apache_beam.utils import shared
@@ -54,6 +58,15 @@
 ModelT = TypeVar('ModelT')
 ExampleT = TypeVar('ExampleT')
 PredictionT = TypeVar('PredictionT')
+_K = TypeVar('_K')
+_INPUT_TYPE = TypeVar('_INPUT_TYPE')
+_OUTPUT_TYPE = TypeVar('_OUTPUT_TYPE')
+
+
+@dataclass
+class PredictionResult:
+  example: Generic[_INPUT_TYPE]
+  inference: Generic[_OUTPUT_TYPE]
 
 
 def _to_milliseconds(time_ns: int) -> int:
@@ -101,12 +114,27 @@ def batch_elements_kwargs(self) -> Mapping[str, Any]:
     return {}
 
 
+@beam.typehints.with_input_types(Union[_INPUT_TYPE, Tuple[_K, _INPUT_TYPE]])
+@beam.typehints.with_output_types(Union[PredictionResult, Tuple[_K, PredictionResult]])  # pylint: disable=line-too-long
 class RunInference(beam.PTransform[beam.PCollection[ExampleT],
                                    beam.PCollection[PredictionT]]):
   """An extensible transform for running inferences.
   Args:
       model_loader: An implementation of ModelLoader.
       clock: A clock implementing get_current_time_in_microseconds.
+
+  This transform takes a PCollection of examples (or features) to be used on
+  an ML model. It will then output inferences (or predictions) for those
+  examples in a PCollection of PredictionResults, containing the input examples
+  and output inferences.
+
+  If examples are paired with keys, it will output a tuple
+  (key, PredictionResult) for each (key, example) input.
+
+  Models for supported frameworks can be loaded via a URI. Supported services
+  can also be used.
+
+  TODO(BEAM-14046): Add and link to help documentation
   """
   def __init__(
       self,

sdks/python/apache_beam/ml/inference/pytorch_inference.py

@@ -27,7 +27,7 @@
 
 import torch
 from apache_beam.io.filesystems import FileSystems
-from apache_beam.ml.inference.api import PredictionResult
+from apache_beam.ml.inference.base import PredictionResult
 from apache_beam.ml.inference.base import InferenceRunner
 from apache_beam.ml.inference.base import ModelLoader
 

sdks/python/apache_beam/ml/inference/pytorch_inference_test.py

@@ -35,7 +35,7 @@
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
 try:
   import torch
-  from apache_beam.ml.inference.api import PredictionResult
+  from apache_beam.ml.inference.base import PredictionResult
   from apache_beam.ml.inference.base import RunInference
   from apache_beam.ml.inference.pytorch_inference import PytorchInferenceRunner
   from apache_beam.ml.inference.pytorch_inference import PytorchModelLoader

sdks/python/apache_beam/ml/inference/sklearn_inference.py

@@ -25,7 +25,7 @@
 from sklearn.base import BaseEstimator
 
 from apache_beam.io.filesystems import FileSystems
-from apache_beam.ml.inference.api import PredictionResult
+from apache_beam.ml.inference.base import PredictionResult
 from apache_beam.ml.inference.base import InferenceRunner
 from apache_beam.ml.inference.base import ModelLoader
 

sdks/python/apache_beam/ml/inference/sklearn_inference_test.py

@@ -30,7 +30,6 @@
 from sklearn import svm
 
 import apache_beam as beam
-from apache_beam.ml.inference import api
 from apache_beam.ml.inference import base
 from apache_beam.ml.inference.sklearn_inference import ModelFileType
 from apache_beam.ml.inference.sklearn_inference import SklearnInferenceRunner
@@ -76,9 +75,9 @@ def test_predict_output(self):
         numpy.array([1, 2, 3]), numpy.array([4, 5, 6]), numpy.array([7, 8, 9])
     ]
     expected_predictions = [
-        api.PredictionResult(numpy.array([1, 2, 3]), 6),
-        api.PredictionResult(numpy.array([4, 5, 6]), 15),
-        api.PredictionResult(numpy.array([7, 8, 9]), 24)
+        base.PredictionResult(numpy.array([1, 2, 3]), 6),
+        base.PredictionResult(numpy.array([4, 5, 6]), 15),
+        base.PredictionResult(numpy.array([7, 8, 9]), 24)
     ]
     inferences = inference_runner.run_inference(batched_examples, fake_model)
     for actual, expected in zip(inferences, expected_predictions):
@@ -126,8 +125,8 @@ def test_pipeline_pickled(self):
       actual = pcoll | base.RunInference(
           SklearnModelLoader(model_uri=temp_file_name))
       expected = [
-          api.PredictionResult(numpy.array([0, 0]), 0),
-          api.PredictionResult(numpy.array([1, 1]), 1)
+          base.PredictionResult(numpy.array([0, 0]), 0),
+          base.PredictionResult(numpy.array([1, 1]), 1)
       ]
       assert_that(
           actual, equal_to(expected, equals_fn=_compare_prediction_result))
@@ -147,8 +146,8 @@ def test_pipeline_joblib(self):
           SklearnModelLoader(
               model_uri=temp_file_name, model_file_type=ModelFileType.JOBLIB))
       expected = [
-          api.PredictionResult(numpy.array([0, 0]), 0),
-          api.PredictionResult(numpy.array([1, 1]), 1)
+          base.PredictionResult(numpy.array([0, 0]), 0),
+          base.PredictionResult(numpy.array([1, 1]), 1)
       ]
       assert_that(
           actual, equal_to(expected, equals_fn=_compare_prediction_result))