| 31 | } |
| 32 | |
| 33 | void OutputModelCoreML( |
| 34 | const TFullModel& model, |
| 35 | const TString& modelFile, |
| 36 | const NJson::TJsonValue& userParameters, |
| 37 | const THashMap<ui32, TString>* catFeaturesHashToString) { |
| 38 | |
| 39 | CoreML::Specification::Model treeModel; |
| 40 | treeModel.set_specificationversion(1); |
| 41 | |
| 42 | auto regressor = treeModel.mutable_treeensembleregressor(); |
| 43 | auto ensemble = regressor->mutable_treeensemble(); |
| 44 | |
| 45 | NCB::NCoreML::TPerTypeFeatureIdxToInputIndex perTypeFeatureIdxToInputIndex; |
| 46 | bool createPipelineModel = model.HasCategoricalFeatures(); |
| 47 | |
| 48 | TString data; |
| 49 | if (createPipelineModel) { |
| 50 | CoreML::Specification::Model pipelineModel; |
| 51 | pipelineModel.set_specificationversion(1); |
| 52 | |
| 53 | auto* container = pipelineModel.mutable_pipeline()->mutable_models(); |
| 54 | NCB::NCoreML::ConfigureCategoricalMappings(model, catFeaturesHashToString, container); |
| 55 | |
| 56 | auto* contained = container->Add(); |
| 57 | auto treeDescription = treeModel.mutable_description(); |
| 58 | NCB::NCoreML::ConfigureTreeModelIO(model, userParameters, regressor, treeDescription, &perTypeFeatureIdxToInputIndex); |
| 59 | |
| 60 | NCB::NCoreML::ConfigureTrees(model, perTypeFeatureIdxToInputIndex, ensemble); |
| 61 | |
| 62 | *contained = treeModel; |
| 63 | |
| 64 | auto pipelineDescription = pipelineModel.mutable_description(); |
| 65 | NCB::NCoreML::ConfigureMetadata(model, userParameters, pipelineDescription); |
| 66 | NCB::NCoreML::ConfigurePipelineModelIO(model, pipelineDescription); |
| 67 | |
| 68 | Y_PROTOBUF_SUPPRESS_NODISCARD pipelineModel.SerializeToString(&data); |
| 69 | } else { |
| 70 | auto description = treeModel.mutable_description(); |
| 71 | NCB::NCoreML::ConfigureMetadata(model, userParameters, description); |
| 72 | NCB::NCoreML::ConfigureTreeModelIO(model, userParameters, regressor, description, &perTypeFeatureIdxToInputIndex); |
| 73 | |
| 74 | NCB::NCoreML::ConfigureTrees(model, perTypeFeatureIdxToInputIndex, ensemble); |
| 75 | |
| 76 | Y_PROTOBUF_SUPPRESS_NODISCARD treeModel.SerializeToString(&data); |
| 77 | } |
| 78 | |
| 79 | TOFStream out(modelFile); |
| 80 | out.Write(data); |
| 81 | } |
| 82 | |
| 83 | void SerializeFullModelToOnnxStream( |
| 84 | const TFullModel& model, |
no test coverage detected