Skip to content

Commit

Permalink
Remove usage of deprecated _serialize (apache#33000)
Browse files Browse the repository at this point in the history
* Remove usage of deprecated _serialize

* Correct assignment

* indentation

* lint

* fmt

* fmt
  • Loading branch information
damccorm authored Nov 1, 2024
1 parent 61268ef commit 90c1ee9
Showing 1 changed file with 9 additions and 1 deletion.
10 changes: 9 additions & 1 deletion sdks/python/apache_beam/ml/inference/onnx_inference.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,7 +116,15 @@ def load_model(self) -> ort.InferenceSession:
# when path is remote, we should first load into memory then deserialize
f = FileSystems.open(self._model_uri, "rb")
model_proto = onnx.load(f)
model_proto_bytes = onnx._serialize(model_proto)
model_proto_bytes = model_proto
if not isinstance(model_proto, bytes):
if (hasattr(model_proto, "SerializeToString") and
callable(model_proto.SerializeToString)):
model_proto_bytes = model_proto.SerializeToString()
else:
raise TypeError(
"No SerializeToString method is detected on loaded model. " +
f"Type of model: {type(model_proto)}")
ort_session = ort.InferenceSession(
model_proto_bytes,
sess_options=self._session_options,
Expand Down

0 comments on commit 90c1ee9

Please sign in to comment.