Funtion to dump trained Scikit-Learn RandomForestRegressor to JSON

rfr_to_json.py

from json import dumps

def rfr_to_json(rfr_object, feature_list, json_filepath=None):
    '''
    Function to convert a scikit-learn RandomForestRegressor object to JSON.
    
    '''

    output_dict = dict()
    output_dict['name'] = 'rf_regression_pipeline'
    output_dict['transformer'] = dict()
    output_dict['transformer']['stages'] = list()

    output_dict['transformer']['stages'].append(dict())
    output_dict['transformer']['stages'][0]['task'] = 'extract-features'
    output_dict['transformer']['stages'][0]['features'] = list()
    
    for i, c in enumerate(feature_list):
        feature_dict = {
            'type': 'double', 
            'name': c, 
            'order': i,
        }
        output_dict['transformer']['stages'][0]['features'].append(feature_dict)
                                                                   
    output_dict['transformer']['stages'].append(dict())
    output_dict['transformer']['stages'][1]['outputType'] = 'double'
    output_dict['transformer']['stages'][1]['pythonObject'] = 'sklearn.ensemble.RandomForestRegressor'
    output_dict['transformer']['stages'][1]['nTrees'] = len(rfr_object.estimators_)
    output_dict['transformer']['stages'][1]['trees'] = list()
                                                               
    for i, model in enumerate(rfr_object.estimators_):
        tree_dict = dict(
            i=i,
            tree_undefined=_tree.TREE_UNDEFINED,
            features=model.tree_.feature.tolist(),
            thresholds=model.tree_.threshold.tolist(),
            children_left=model.tree_.children_left.tolist(),
            children_right=model.tree_.children_right.tolist(),
            values=[val[0][0] for val in model.tree_.value]
        )
        output_dict['transformer']['stages'][1]['trees'].append(tree_dict)

    if json_filepath is not None:
        with open(json_filepath, 'w') as json_file:
            dump(output_dict, json_file)
    else:
        print(dumps(output_dict))