[Python][ML] Refactor batch conversion into helpers

Author: siliataider

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py

@@ -345,35 +345,29 @@ def GetSample(self):
             np.zeros((self.batch_size)).reshape(-1, 1),
         )
 
-    def ConvertBatchToNumpy(self, batch) -> np.ndarray:
-        """Convert a RTensor into a NumPy array
-
-        Args:
-            batch (RTensor): Batch returned from the DataLoader
-
-        Returns:
-            np.ndarray: converted batch
-        """
+    def _get_raw_array(self, batch) -> np.ndarray:
         try:
             import numpy as np
         except ImportError:
             raise ImportError("Failed to import numpy needed for the ML dataloader")
 
         data = batch.GetData()
         batch_size, num_columns = tuple(batch.GetShape())
-
         data.reshape((batch_size * num_columns,))
 
-        return_data = np.asarray(data).reshape(batch_size, num_columns)
+        return np.asarray(data).reshape(batch_size, num_columns)
 
+    def _split_target_and_weights(
+        self, data: np.ndarray
+    ) -> np.ndarray | Tuple[np.ndarray, np.ndarray] | Tuple[np.ndarray, np.ndarray, np.ndarray]:
         # Splice target column from the data if target is given
         if self.target_given:
-            train_data = return_data[:, self.train_indices]
-            target_data = return_data[:, self.target_indices]
+            train_data = data[:, self.train_indices]
+            target_data = data[:, self.target_indices]
 
             # Splice weight column from the data if weight is given
             if self.weights_given:
-                weights_data = return_data[:, self.weights_index]
+                weights_data = data[:, self.weights_index]
 
                 if len(self.target_indices) == 1:
                     return train_data, target_data.reshape(-1, 1), weights_data.reshape(-1, 1)
@@ -385,7 +379,18 @@ def ConvertBatchToNumpy(self, batch) -> np.ndarray:
 
             return train_data, target_data
 
-        return return_data
+        return data
+
+    def ConvertBatchToNumpy(self, batch) -> np.ndarray:
+        """Convert a RTensor into a NumPy array
+
+        Args:
+            batch (RTensor): Batch returned from the DataLoader
+
+        Returns:
+            np.ndarray: converted batch
+        """
+        return self._split_target_and_weights(self._get_raw_array(batch))
 
     def ConvertBatchToPyTorch(self, batch: Any, device=None) -> torch.Tensor:
         """Convert a RTensor into a PyTorch tensor
@@ -396,36 +401,14 @@ def ConvertBatchToPyTorch(self, batch: Any, device=None) -> torch.Tensor:
         Returns:
             torch.Tensor: converted batch
         """
-        import numpy as np
         import torch
 
-        data = batch.GetData()
-        batch_size, num_columns = tuple(batch.GetShape())
-
-        data.reshape((batch_size * num_columns,))
-
-        return_data = torch.as_tensor(np.asarray(data), device=device).reshape(batch_size, num_columns)
-
-        # Splice target column from the data if target is given
-        if self.target_given:
-            train_data = return_data[:, self.train_indices]
-            target_data = return_data[:, self.target_indices]
-
-            # Splice weight column from the data if weight is given
-            if self.weights_given:
-                weights_data = return_data[:, self.weights_index]
-
-                if len(self.target_indices) == 1:
-                    return train_data, target_data.reshape(-1, 1), weights_data.reshape(-1, 1)
-
-                return train_data, target_data, weights_data.reshape(-1, 1)
-
-            if len(self.target_indices) == 1:
-                return train_data, target_data.reshape(-1, 1)
-
-            return train_data, target_data
-
-        return return_data
+        split = self._split_target_and_weights(self._get_raw_array(batch))
+        return (
+            tuple(torch.as_tensor(arr, device=device) for arr in split)
+            if isinstance(split, tuple)
+            else torch.as_tensor(split, device=device)
+        )
 
     def ConvertBatchToTF(self, batch: Any) -> Any:
         """
@@ -439,12 +422,9 @@ def ConvertBatchToTF(self, batch: Any) -> Any:
         """
         import tensorflow as tf
 
-        data = batch.GetData()
-        batch_size, num_columns = tuple(batch.GetShape())
-
-        data.reshape((batch_size * num_columns,))
-
-        return_data = tf.constant(data, shape=(batch_size, num_columns))
+        arr = self._get_raw_array(batch)
+        batch_size = arr.shape[0]
+        return_data = tf.constant(arr)
 
         if batch_size != self.batch_size:
             return_data = tf.pad(return_data, tf.constant([[0, self.batch_size - batch_size], [0, 0]]))
@@ -464,6 +444,7 @@ def ConvertBatchToTF(self, batch: Any) -> Any:
 
         return return_data
 
+
     # Return a batch when available
     def GetTrainBatch(self) -> Any:
         """Return the next training batch of data from the given RDataFrame
@@ -717,6 +698,11 @@ def as_torch(self, device: str | torch.device | None = None) -> FormattedLoader:
         Args:
             device: If given, the returned tensors are moved to the specified device.
         """
+        try:
+            import torch  # noqa F401
+        except ImportError:
+            raise ImportError("Failed to import torch needed for the ML dataloader")
+
         self._ensure_created()
         conversion_fn = lambda batch: self._internal.ConvertBatchToPyTorch(batch, device)  # noqa: E731
         return FormattedLoader(self._internal, conversion_fn, self._is_training)
@@ -726,7 +712,10 @@ def as_tensorflow(self) -> tf.data.Dataset:
         \ingroup Py_ML
         Return a tf.data.Dataset over batches as TensorFlow tensors.
         """
-        import tensorflow as tf
+        try:
+            import tensorflow as tf
+        except ImportError:
+            raise ImportError("Failed to import tensorflow needed for the ML dataloader")
 
         self._ensure_created()