<commit_before>import functools

from federatedml.ensemble.basic_algorithms import HeteroDecisionTreeGuest
from federatedml.ensemble.boosting.hetero import hetero_fast_secureboost_plan as plan
from federatedml.util import consts

from federatedml.ensemble.basic_algorithms.decision_tree.tree_core.node import Node

from arch.api.utils import log_utils

LOGGER = log_utils.getLogger()


class HeteroFastDecisionTreeGuest(HeteroDecisionTreeGuest):

    def __init__(self, tree_param):
        super(HeteroFastDecisionTreeGuest, self).__init__(tree_param)
        self.node_plan = []
        self.node_plan_idx = 0
        self.tree_type = None
        self.target_host_id = -1
        self.guest_depth = 0
        self.host_depth = 0
        self.cur_dep = 0
        self.use_guest_feat_when_predict = False

    def use_guest_feat_only_predict_mode(self):
        self.use_guest_feat_when_predict = True

    def set_tree_work_mode(self, tree_type, target_host_id):
        self.tree_type, self.target_host_id = tree_type, target_host_id

    def set_layered_depth(self, guest_depth, host_depth):
        self.guest_depth, self.host_depth = guest_depth, host_depth

    def initialize_node_plan(self):
        if self.tree_type == plan.tree_type_dict['layered_tree']:
            self.node_plan = plan.create_layered_tree_node_plan(guest_depth=self.guest_depth,
                                                                host_depth=self.host_depth,
                                                                host_list=self.host_party_idlist)
            self.max_depth = len(self.node_plan)
            LOGGER.debug('max depth reset to {}, cur node plan is {}'.format(self.max_depth, self.node_plan))
        else:
            self.node_plan = plan.create_node_plan(self.tree_type, self.target_host_id, self.max_depth)

    def get_node_plan(self, idx):
        return self.node_plan[idx]

    def host_id_to_idx(self, host_id):
        if host_id == -1:
            return -1
        return self.host_party_idlist.index(host_id)

    def compute_best_splits_with_node_plan(self, tree_action, target_host_idx, node_map: dict, dep: int,
                                           batch_idx: int, mode=consts.MIX_TREE):

        LOGGER.debug('node plan at dep {} is {}'.format(dep, (tree_action, target_host_idx)))

        cur_best_split = []

        if tree_action == plan.tree_actions['guest_only']:
            acc_histograms = self.get_local_histograms(node_map, ret='tensor')
            cur_best_split = self.splitter.find_split(acc_histograms, self.valid_features,
                                                      self.data_bin._partitions, self.sitename,
                                                      self.use_missing, self.zero_as_missing)
            LOGGER.debug('computing local splits done')

        if tree_action == plan.tree_actions['host_only']:

            self.federated_find_split(dep, batch_idx, idx=target_host_idx)

            if mode == consts.LAYERED_TREE:
                host_split_info = self.sync_final_split_host(dep, batch_idx, idx=target_host_idx)
                LOGGER.debug('get encrypted split value from host')

                cur_best_split = self.merge_splitinfo(splitinfo_guest=[],
                                                      splitinfo_host=host_split_info,
                                                      merge_host_split_only=True)

        return cur_best_split

    def assign_instances_to_new_node_with_node_plan(self, dep, tree_action, mode=consts.MIX_TREE):

        LOGGER.info("redispatch node of depth {}".format(dep))
        dispatch_node_method = functools.partial(self.assign_a_instance,
                                                 tree_=self.tree_node,
                                                 decoder=self.decode,
                                                 sitename=self.sitename,
                                                 split_maskdict=self.split_maskdict,
                                                 bin_sparse_points=self.bin_sparse_points,<commit_msg>_partitions -> partitions hetero_fast_decision_tree_guest.py<commit_after>import functools

from federatedml.ensemble.basic_algorithms import HeteroDecisionTreeGuest
from federatedml.ensemble.boosting.hetero import hetero_fast_secureboost_plan as plan
from federatedml.util import consts

from federatedml.ensemble.basic_algorithms.decision_tree.tree_core.node import Node

from arch.api.utils import log_utils

LOGGER = log_utils.getLogger()


class HeteroFastDecisionTreeGuest(HeteroDecisionTreeGuest):

    def __init__(self, tree_param):
        super(HeteroFastDecisionTreeGuest, self).__init__(tree_param)
        self.node_plan = []
        self.node_plan_idx = 0
        self.tree_type = None
        self.target_host_id = -1
        self.guest_depth = 0
        self.host_depth = 0
        self.cur_dep = 0
        self.use_guest_feat_when_predict = False

    def use_guest_feat_only_predict_mode(self):
        self.use_guest_feat_when_predict = True

    def set_tree_work_mode(self, tree_type, target_host_id):
        self.tree_type, self.target_host_id = tree_type, target_host_id

    def set_layered_depth(self, guest_depth, host_depth):
        self.guest_depth, self.host_depth = guest_depth, host_depth

    def initialize_node_plan(self):
        if self.tree_type == plan.tree_type_dict['layered_tree']:
            self.node_plan = plan.create_layered_tree_node_plan(guest_depth=self.guest_depth,
                                                                host_depth=self.host_depth,
                                                                host_list=self.host_party_idlist)
            self.max_depth = len(self.node_plan)
            LOGGER.debug('max depth reset to {}, cur node plan is {}'.format(self.max_depth, self.node_plan))
        else:
            self.node_plan = plan.create_node_plan(self.tree_type, self.target_host_id, self.max_depth)

    def get_node_plan(self, idx):
        return self.node_plan[idx]

    def host_id_to_idx(self, host_id):
        if host_id == -1:
            return -1
        return self.host_party_idlist.index(host_id)

    def compute_best_splits_with_node_plan(self, tree_action, target_host_idx, node_map: dict, dep: int,
                                           batch_idx: int, mode=consts.MIX_TREE):

        LOGGER.debug('node plan at dep {} is {}'.format(dep, (tree_action, target_host_idx)))

        cur_best_split = []

        if tree_action == plan.tree_actions['guest_only']:
            acc_histograms = self.get_local_histograms(node_map, ret='tensor')
            cur_best_split = self.splitter.find_split(acc_histograms, self.valid_features,
                                                      self.data_bin.partitions, self.sitename,
                                                      self.use_missing, self.zero_as_missing)
            LOGGER.debug('computing local splits done')

        if tree_action == plan.tree_actions['host_only']:

            self.federated_find_split(dep, batch_idx, idx=target_host_idx)

            if mode == consts.LAYERED_TREE:
                host_split_info = self.sync_final_split_host(dep, batch_idx, idx=target_host_idx)
                LOGGER.debug('get encrypted split value from host')

                cur_best_split = self.merge_splitinfo(splitinfo_guest=[],
                                                      splitinfo_host=host_split_info,
                                                      merge_host_split_only=True)

        return cur_best_split

    def assign_instances_to_new_node_with_node_plan(self, dep, tree_action, mode=consts.MIX_TREE):

        LOGGER.info("redispatch node of depth {}".format(dep))
        dispatch_node_method = functools.partial(self.assign_a_instance,
                                                 tree_=self.tree_node,
                                                 decoder=self.decode,
                                                 sitename=self.sitename,
                                                 split_maskdict=self.split_maskdict,
                                                 bin_sparse_points=self.bin_sparse_points,
{'fraction_non_alphanumeric': 0.0741236611489776, 'fraction_numerical': 0.0017039922103213243, 'mean_word_length': 2.0343426883308715, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 48, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '381451', 'n_tokens_mistral': 2380, 'n_tokens_neox': 2302, 'n_words': 377}