Method ConfigureTreeModelIO

catboost/libs/model/model_export/coreml_helpers.cpp:227–289

}

void NCB::NCoreML::ConfigureTreeModelIO(
    const TFullModel& model,
    const NJson::TJsonValue& userParameters,
    TreeEnsembleRegressor* regressor,
    ModelDescription* description,
    TPerTypeFeatureIdxToInputIndex* perTypeFeatureIdxToInputIndex) {

    ConfigureFloatInput(model, description, &(perTypeFeatureIdxToInputIndex->ForFloatFeatures));

    size_t catFeaturesCount  = model.ModelTrees->GetCatFeatures().size();
    TVector<int> categoricalFlatIndexes(catFeaturesCount);
    for (const auto& catFeature: model.ModelTrees->GetCatFeatures()) {
        categoricalFlatIndexes[catFeature.Position.Index] = catFeature.Position.FlatIndex;
    }

    for (const auto& oneHotFeature : model.ModelTrees->GetOneHotFeatures()) {
        (*perTypeFeatureIdxToInputIndex).ForCatFeatures[oneHotFeature.CatFeatureIndex] = description->input().size();

        auto feature = description->add_input();
        int flatFeatureIndex = categoricalFlatIndexes[oneHotFeature.CatFeatureIndex];

        feature->set_name(("mapped_feature_" + std::to_string(flatFeatureIndex)).c_str());

        auto featureType = new FeatureType();
        featureType->set_isoptional(false);
        featureType->set_allocated_int64type(new Int64FeatureType());
        feature->set_allocated_type(featureType);
    }

    const auto classesCount = static_cast<size_t>(model.ModelTrees->GetDimensionsCount());
    regressor->mutable_treeensemble()->set_numpredictiondimensions(classesCount);
    if (classesCount == 1) {
        regressor->mutable_treeensemble()->add_basepredictionvalue(
            model.ModelTrees->GetScaleAndBias().GetOneDimensionalBias(
                "Non single-dimension approxes are not supported")
        );
    } else {
        for (size_t outputIdx = 0; outputIdx < classesCount; ++outputIdx) {
            regressor->mutable_treeensemble()->add_basepredictionvalue(0.0);
        }
    }

    auto outputPrediction = description->add_output();
    outputPrediction->set_name("prediction");
    description->set_predictedfeaturename("prediction");
    description->set_predictedprobabilitiesname("prediction");

    auto featureType = outputPrediction->mutable_type();
    featureType->set_isoptional(false);

    auto outputArray = new ArrayFeatureType();
    outputArray->set_datatype(ArrayFeatureType::DOUBLE);
    outputArray->add_shape(classesCount);

    featureType->set_allocated_multiarraytype(outputArray);

    const auto& prediction_type = userParameters["prediction_type"].GetString();
    if (prediction_type == "probability") {