damccorm commented on code in PR #28243:
URL: https://github.com/apache/beam/pull/28243#discussion_r1310799650
##########
website/www/site/content/en/documentation/sdks/python-machine-learning.md:
##########
@@ -215,6 +215,26 @@ with pipeline as p:
If you are unsure if your data is keyed, you can also use
`MaybeKeyedModelHandler`.
+You can also use a `KeyedModelHandler` to load several different models based
on their associated key:
+
+```
+from apache_beam.ml.inference.base import KeyedModelHandler
+keyed_model_handler = KeyedModelHandler([
+ KeyModelMapping(['key1'], PytorchModelHandlerTensor(<config1>)),
+ KeyModelMapping(['key2', 'key3'], PytorchModelHandlerTensor(<config2>))
+])
+with pipeline as p:
+ data = p | beam.Create([
+ ('key1', torch.tensor([[1,2,3],[4,5,6],...])),
+ ('key2', torch.tensor([[1,2,3],[4,5,6],...])),
+ ('key3', torch.tensor([[1,2,3],[4,5,6],...])),
+ ])
+ predictions = data | RunInference(keyed_model_handler)
+```
+
+The previous example will load a model using `config1` and use that for
inference for all examples associated
+with `key1`, and will load a model using `config2` and use that for all
examples associated with `key2` and `key3`.
Review Comment:
TODO - this still needs to talk about memory pitfalls/memory management
--
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]
