fix: compute first size

Author: cuongth95

ditec_wdn_dataset/core/datasets_large.py

@@ -171,6 +171,8 @@ def compute_first_size(self) -> int:
             my_zarr = self.get_array_by_key(node_attr)
             num_samples = my_zarr.shape[0]
             assert num_samples is not None and isinstance(num_samples, int)
+            if num_samples != self.attrs["num_samples"]:
+                print(f"WARN! The number of samples ({num_samples}) does not match the #samples ({self.attrs['num_samples']}) recorded in metadata ")
             return num_samples
 
     def __init__(

ditec_wdn_dataset/hf/dataset.py

@@ -277,12 +277,14 @@ def get_array_by_key(self, key: str) -> hfd.Dataset:
         def compute_first_size(self) -> int:
             if len(self.sorted_node_attrs) <= 0:
                 return 0
-            num_samples = self.root.num_rows
-            if isinstance(num_samples, dict):
-                total_samples = sum(num_samples.values())
-                return total_samples
-            else:
-                return num_samples
+            # num_samples = self.root.num_rows
+            # if isinstance(num_samples, dict):
+            #     total_samples = sum(num_samples.values())
+            #     return total_samples
+            # else:
+            #     return num_samples
+
+            return self.attrs["num_samples"]
 
     def __init__(
         self,