(beam) 01/01: Fix pre/postprocess type hints

[da...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

damccorm pushed a commit to branch users/damccorm/th
in repository https://gitbox.apache.org/repos/asf/beam.git

commit a4f3165fb07430fd7fc28862d53a1ed855951f94
Author: Danny McCormick <da...@google.com>
AuthorDate: Fri Mar 14 14:08:30 2025 -0400

    Fix pre/postprocess type hints
---
 sdks/python/apache_beam/ml/inference/base.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/sdks/python/apache_beam/ml/inference/base.py b/sdks/python/apache_beam/ml/inference/base.py
index 5cc38cd842c..e4c3e4cab5e 100644
--- a/sdks/python/apache_beam/ml/inference/base.py
+++ b/sdks/python/apache_beam/ml/inference/base.py
@@ -272,7 +272,7 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
 
   def with_preprocess_fn(
       self, fn: Callable[[PreProcessT], ExampleT]
-  ) -> 'ModelHandler[PreProcessT, PredictionT, ModelT, PreProcessT]':
+  ) -> 'ModelHandler[PreProcessT, PredictionT, ModelT]':
     """Returns a new ModelHandler with a preprocessing function
     associated with it. The preprocessing function will be run
     before batching/inference and should map your input PCollection
@@ -283,7 +283,7 @@ class ModelHandler(Generic[ExampleT, PredictionT, ModelT]):
 
   def with_postprocess_fn(
       self, fn: Callable[[PredictionT], PostProcessT]
-  ) -> 'ModelHandler[ExampleT, PostProcessT, ModelT, PostProcessT]':
+  ) -> 'ModelHandler[ExampleT, PostProcessT, ModelT]':
     """Returns a new ModelHandler with a postprocessing function
     associated with it. The postprocessing function will be run
     after inference and should map the base ModelHandler's output