code: def sample_per_label(
      train_data, sample_count=None, random_state=None
  ):
      train_data_binarized_label = train_data.label.get_dummies()
      labels = [
          int(label_name.split("_")[1])
          for label_name in train_data_binarized_label.columns
      ]
      sampled_data_df = []
      sampled_n = 0
      data_n = train_data.shape[0]
      sample_n_per_label = {}
      for i, label in enumerate(labels):
          label_data = train_data.iloc(train_data.label == label)
          if i == len(labels) - 1:
              # last label:
              to_sample_n = sample_count - sampled_n
          else:
              to_sample_n = round(
                  label_data.shape[0] / data_n * sample_count
              )
          label_sampled_data = sample_data(
              df=label_data, n=to_sample_n, random_state=random_state
          )
          if label_sampled_data:
              sample_n_per_label[label] = label_sampled_data.shape[0]
          else:
              sample_n_per_label[label] = 0
          if label_sampled_data is not None:
              sampled_data_df.append(label_sampled_data)
              sampled_n += label_sampled_data.shape[0]
      sampled_data = None
      if sampled_data_df:
          sampled_data = DataFrame.vstack(sampled_data_df)
      return sampled_data, sample_n_per_label, labels
  
file_path: FATE/python/fate/ml/model_selection/data_split.py
imports: import logging
  from fate.arch import Context
  from fate.arch.dataframe import DataFrame
  from ..abc.module import Module
label_desc: Use image node_id 2 for calling a global function with example usage: sample_per_label(train_data, sample_count, random_state) and returns: sampled_data,  sample_n_per_label,  labels