[Python] Optimize RunInference PyTorch/HF get_num_bytes and Sklearn Pandas batching

sdks/python/apache_beam/ml/inference/huggingface_inference.py

@@ -404,7 +404,8 @@ def get_num_bytes(
       return sum(sys.getsizeof(element) for element in batch)
     else:
       return sum(
-          (el.element_size() for tensor in batch for el in tensor.values()))
+          el.element_size() * el.nelement() for tensor in batch
+          for el in tensor.values())
 
   def get_metrics_namespace(self) -> str:
     """
@@ -607,9 +608,12 @@ def get_num_bytes(
     """
     if self._framework == "tf":
       return sum(sys.getsizeof(element) for element in batch)
-    else:
+    elif batch and isinstance(batch[0], dict):
       return sum(
-          (el.element_size() for tensor in batch for el in tensor.values()))
+          el.element_size() * el.nelement() for tensor in batch
+          for el in tensor.values())
+    else:
+      return sum(tensor.element_size() * tensor.nelement() for tensor in batch)
 
   def get_metrics_namespace(self) -> str:
     """

sdks/python/apache_beam/ml/inference/pytorch_inference.py

@@ -360,7 +360,7 @@ def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int:
     Returns:
       The number of bytes of data for a batch of Tensors.
     """
-    return sum((el.element_size() for tensor in batch for el in tensor))
+    return sum(tensor.element_size() * tensor.nelement() for tensor in batch)
 
   def get_metrics_namespace(self) -> str:
     """
@@ -609,7 +609,8 @@ def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int:
     """
     # If elements in `batch` are provided as a dictionaries from key to Tensors
     return sum(
-        (el.element_size() for tensor in batch for el in tensor.values()))
+        el.element_size() * el.nelement() for tensor in batch
+        for el in tensor.values())
 
   def get_metrics_namespace(self) -> str:
     """

sdks/python/apache_beam/ml/inference/sklearn_inference.py

@@ -212,6 +212,8 @@ def _default_pandas_inference_fn(
   # vectorize data for better performance
   vectorized_batch = pandas.concat(batch, axis=0)
   predictions = model.predict(vectorized_batch)
+  if len(batch) == vectorized_batch.shape[0]:
+    return predictions, batch
   splits = [
       vectorized_batch.iloc[[i]] for i in range(vectorized_batch.shape[0])
   ]