yeandy commented on code in PR #23266:
URL: https://github.com/apache/beam/pull/23266#discussion_r973098948
##########
sdks/python/apache_beam/ml/inference/pytorch_inference.py:
##########
@@ -277,6 +262,8 @@ def run_inference(
inference_args: Non-batchable arguments required as inputs to the model's
forward() function. Unlike Tensors in `batch`, these parameters will
not be dynamically batched
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/sklearn_inference.py:
##########
@@ -124,6 +107,7 @@ def run_inference(
model: A numpy model or pipeline. Must implement predict(X).
Where the parameter X is a numpy array.
inference_args: Any additional arguments for an inference.
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -92,14 +110,16 @@ def run_inference(
self,
batch: Sequence[ExampleT],
model: ModelT,
- inference_args: Optional[Dict[str, Any]] = None) ->
Iterable[PredictionT]:
+ inference_args: Optional[Dict[str, Any]] = None,
+ drop_example: bool = False) -> Iterable[PredictionT]:
"""Runs inferences on a batch of examples.
Args:
batch: A sequence of examples or features.
model: The model used to make inferences.
inference_args: Extra arguments for models whose inference call requires
extra parameters.
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/base_test.py:
##########
@@ -251,6 +281,19 @@ def
test_run_inference_keyed_examples_with_unkeyed_model_handler(self):
| 'RunKeyed' >> base.RunInference(model_handler))
pipeline.run()
+ def test_drop_example_prediction_result(self):
+ def assert_drop_example(prediction_result):
+ assert prediction_result.example is None
+
+ pipeline = TestPipeline()
+ examples = [1, 3, 5]
+ model_handler = FakeModelHandlerWithPredictionResult()
Review Comment:
Shouldn't this have `drop_example=True`?
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -385,19 +413,22 @@ def __init__(
self,
model_handler: ModelHandler[ExampleT, PredictionT, Any],
clock,
- metrics_namespace):
+ metrics_namespace,
+ drop_example=False):
"""A DoFn implementation generic to frameworks.
Args:
model_handler: An implementation of ModelHandler.
clock: A clock implementing time_ns. *Used for unit testing.*
metrics_namespace: Namespace of the transform to collect metrics.
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/base_test.py:
##########
@@ -48,13 +48,40 @@ def run_inference(
self,
batch: Sequence[int],
model: FakeModel,
- inference_args=None) -> Iterable[int]:
+ inference_args=None,
+ drop_example=False) -> Iterable[int]:
if self._fake_clock:
self._fake_clock.current_time_ns += 3_000_000 # 3 milliseconds
for example in batch:
yield model.predict(example)
+class FakeModelHandlerWithPredictionResult(base.ModelHandler[int,
Review Comment:
The name of this class was initially confusing to me. I think I see what you
mean since `run_inference` returns `PredictionResult`, but I'm wondering if
there's a better way to word this. Maybe
`FakeModelHandlerReturnsPredictionResult`?
##########
sdks/python/apache_beam/ml/inference/sklearn_inference.py:
##########
@@ -194,6 +178,7 @@ def run_inference(
model: A dataframe model or pipeline. Must implement predict(X).
Where the parameter X is a pandas dataframe.
inference_args: Any additional arguments for an inference.
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/base.py:
##########
@@ -291,11 +314,13 @@ def __init__(
inference_args: Extra arguments for models whose inference call
requires
extra parameters.
metrics_namespace: Namespace of the transform to collect metrics.
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
##########
sdks/python/apache_beam/ml/inference/pytorch_inference.py:
##########
@@ -170,6 +153,7 @@ def run_inference(
inference_args: Non-batchable arguments required as inputs to the model's
forward() function. Unlike Tensors in `batch`, these parameters will
not be dynamically batched
+ drop_example: Enable this to drop the example from PredictionResult
Review Comment:
```suggestion
drop_example: Boolean flag indicating whether or not to drop the
example from PredictionResult
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
