def get_max_available_num(all_datasets, dataset_weight):
    """
    Get the maximum number of samples for each dataset.
    """
    max_available_num = {}  # the maximum number of samples each dataset can provide

    for key in dataset_weight.keys():  # assume this key is 100% used, O(n^2) time complexity
        this_available_num = {key: len(all_datasets[key])}

        for check_key in dataset_weight.keys():  # check if other key can provide enough number of batches
            if key == check_key:
                continue
            if int(len(all_datasets[key]) * (dataset_weight[check_key] / dataset_weight[key])) > len(all_datasets[check_key]):  # the check_key cannot provide enough data
                this_available_num = None  # not satisfied
                break
            else:  # this check_key is satisfied
                this_available_num[check_key] = int(len(all_datasets[key]) * (dataset_weight[check_key] / dataset_weight[key]))

        if this_available_num is not None:
            if len(max_available_num) == 0:
                max_available_num = this_available_num
            else:
                for key2 in dataset_weight.keys():  # update for all keys
                    if this_available_num[key2] > max_available_num[key2]:
                        max_available_num[key2] = this_available_num[key2]

    return max_available_num

    Get the maximum number of samples for each dataset.
    
smoe/entrypoint/expert_construction/llama_get_hidden_features.py