[damccorm]

Description

Optimizes two hot paths in apache_beam.ml.inference:

1. PytorchModelHandlerTensor, PytorchModelHandlerKeyedTensor, HuggingFaceModelHandlerTensor, and HuggingFaceModelHandlerKeyedTensor (get_num_bytes):
   • Previously, get_num_bytes iterated over tensor elements in Python (sum(el.element_size() for tensor in batch for el in tensor)), which both iterated over the first dimension in Python and undercounted multi-dimensional tensors (summing element_size() per row rather than total elements).
   • Replacing with tensor.element_size() * tensor.nelement() computes the exact tensor byte footprint in $O(1)$ per tensor without iterating tensor slices in Python (up to 2,028x faster on 256x256 float32 tensor batches: 4.17s -> 0.0021s for 1,000 tensors).
2. _default_pandas_inference_fn in sklearn_inference.py:
   • When every input DataFrame in batch is a single row (len(batch) == vectorized_batch.shape[0]), returns (predictions, batch) directly instead of slicing vectorized_batch.iloc[[i]] $N$ times (35.65x faster on single-row DataFrame batches).

Testing

• Verified with pytest sdks/python/apache_beam/ml/inference/pytorch_inference_test.py sdks/python/apache_beam/ml/inference/sklearn_inference_test.py sdks/python/apache_beam/ml/inference/huggingface_inference_test.py (64 passed).

[damccorm]

/gemini review

[github-actions]

📋 Review Summary

This PR significantly improves the performance and accuracy of get_num_bytes for PyTorch and HuggingFace tensor model handlers by replacing element-wise Python iteration with tensor.element_size() * tensor.nelement(), and optimizes _default_pandas_inference_fn in scikit-learn inference by avoiding redundant iloc slicing when inputs are single-row DataFrames.

🔍 General Feedback

• The changes are clean, accurate, and well-targeted at hot paths during batch inference.
• Fixing HuggingFaceModelHandlerTensor.get_num_bytes to handle non-dict PyTorch tensors also resolves an existing bug where calling .values() on a torch.Tensor would fail or reference the bound method rather than computing tensor size.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.