yeandy commented on code in PR #23642:
URL: https://github.com/apache/beam/pull/23642#discussion_r996030825
##########
sdks/python/apache_beam/ml/inference/sklearn_inference.py:
##########
@@ -83,14 +83,21 @@ def _convert_to_result(
]
return [PredictionResult(x, y) for x, y in zip(batch, predictions)]
+def _default_numpy_inference_fn(
+ model: BaseEstimator, batch: Sequence[numpy.ndarray], inference_args:
Optional[Dict[str, Any]] = None
+) -> Any:
+ # vectorize data for better performance
+ vectorized_batch = numpy.stack(batch, axis=0)
+ return model.predict(vectorized_batch)
class SklearnModelHandlerNumpy(ModelHandler[numpy.ndarray,
PredictionResult,
BaseEstimator]):
def __init__(
self,
model_uri: str,
- model_file_type: ModelFileType = ModelFileType.PICKLE):
+ model_file_type: ModelFileType = ModelFileType.PICKLE,
+ inference_fn: Optional[Callable[[BaseEstimator, Sequence[numpy.ndarray],
Optional[Dict[str, Any]]], Any]] = _default_numpy_inference_fn):
Review Comment:
This is getting unwieldy. Wondering if there's a better way to do this.
Could we at least refactor the type for `inference_args: Optional[Dict[str,
Any]]`?
##########
sdks/python/apache_beam/ml/inference/sklearn_inference.py:
##########
@@ -18,7 +18,7 @@
import enum
import pickle
import sys
-from typing import Any
+from typing import Any, Callable
from typing import Dict
Review Comment:
Let's import `Callable` in a separate line
```suggestion
from typing import Any
form typing import Callable
from typing import Dict
```
--
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]
