解析lightgbm的txt模型文件

根据近期的github方案，实现对txt格式的pmml文件的加载

添加依赖

org.jpmml
pmml-lightgbm
1.5.4


org.jpmml
pmml-evaluator
1.6.6


org.jpmml
pmml-model
1.6.6

工具类

import lombok.extern.slf4j.Slf4j;
import org.dmg.pmml.PMML;
import org.jpmml.evaluator.FieldValue;
import org.jpmml.evaluator.InputField;
import org.jpmml.evaluator.ModelEvaluator;
import org.jpmml.evaluator.ModelEvaluatorBuilder;
import org.jpmml.evaluator.ModelEvaluatorFactory;
import org.jpmml.evaluator.ProbabilityDistribution;
import org.jpmml.evaluator.TargetField;
import org.jpmml.lightgbm.GBDT;
import org.jpmml.lightgbm.HasLightGBMOptions;
import org.jpmml.lightgbm.LightGBMUtil;
import org.jpmml.model.metro.MetroJAXBUtil;
import org.springframework.core.io.ClassPathResource;
import org.springframework.core.io.Resource;

import java.io.File;
import java.io.FileOutputStream;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
* 加载、初始化 PMML模型文件 ：
* 依赖 pmml-lightgbm-1.5.0（AGPL-3.0 License）
*

* 解析PMML文件 @link https://github.com/jpmml/jpmml-lightgbm
* 生成evaluator @link https://github.com/jpmml/jpmml-evaluator
*/
@Slf4j
public class LightgbmTxtInitializer {

// description = "Custom objective function"
private static String objectiveFunction = null;
// description = "Transform LightGBM-style trees to PMML-style trees",
private static boolean compact = true;
// description = "Treat Not-a-Number (NaN) values as missing values",
private static boolean nanAsMissing = true;
// description = "Limit the number of trees. Defaults to all trees"
private static Integer numIteration = null;
// description = "Target name. Defaults to "_target""
private static String targetName = null;
// description = "Target categories. Defaults to 0-based index [0, 1, .., num_class - 1]"
private static List targetCategories = null;

public static void main(String[] output) throws Exception {
Resource resource = new ClassPathResource("lightgbm_model.txt");
InputStream pmmlFileInputStream = resource.getInputStream();
// 生成模型执行器
ModelEvaluator evaluator = initEvaluator(pmmlFileInputStream);
// 打印特征参数
List inputFields = evaluator.getInputFields();
log.info("ModelEvaluator featureNames:" + inputFields);
// 调试执行预测
Map<String, Number> waitPreSample = new HashMap<>(8);
waitPreSample.put("0", 0.1);
waitPreSample.put("1", 0.2);
waitPreSample.put("2", 0.3);
String predictedValue = getPredictedValue(waitPreSample, evaluator);

pmmlFileInputStream.close();
}

public static ModelEvaluator initEvaluator(InputStream pmmlFileInputStream) throws Exception {
GBDT gbdt;
long begin = System.currentTimeMillis();

gbdt = LightGBMUtil.loadGBDT(pmmlFileInputStream);
log.info("Loaded GBDT in {} ms.", (System.currentTimeMillis() - begin));

if (objectiveFunction != null) {
log.info("Setting custom objective function");
gbdt.setObjectiveFunction(LightGBMUtil.parseObjectiveFunction(objectiveFunction));
}
Map<String, Object> options = new LinkedHashMap<>();
options.put(HasLightGBMOptions.OPTION_COMPACT, compact);
options.put(HasLightGBMOptions.OPTION_NAN_AS_MISSING, nanAsMissing);
options.put(HasLightGBMOptions.OPTION_NUM_ITERATION, numIteration);

// 生成标准PMML
begin = System.currentTimeMillis();
PMML pmml;
pmml = gbdt.encodePMML(options, targetName, targetCategories);
long end = System.currentTimeMillis();
log.info("Converted GBDT to PMML in {} ms.", (System.currentTimeMillis() - begin));

// no need
// 输出PMML格式文件
begin = System.currentTimeMillis();
File outputFile = new File("E://t.pmml");
OutputStream os = new FileOutputStream(outputFile);
MetroJAXBUtil.marshalPMML(pmml, os);
log.info("Marshalled PMML in {} ms.", (System.currentTimeMillis() - begin));

// 生成evaluator
begin = System.currentTimeMillis();
ModelEvaluatorFactory modelEvaluatorFactory = ModelEvaluatorFactory.newInstance();
ModelEvaluatorBuilder modelEvaluatorBuilder = new ModelEvaluatorBuilder(pmml);
modelEvaluatorBuilder.setModelEvaluatorFactory(modelEvaluatorFactory);
ModelEvaluator<?> evaluator = modelEvaluatorBuilder.build();
evaluator.verify();
log.info("Init evaluator in {} ms.", (System.currentTimeMillis() - begin));
return evaluator;
}

public static String getPredictedValue(Map<String, ?> argumentMap,
ModelEvaluator evaluator) { // 预测计算 Map evaluateResult = evaluator.evaluate(argumentMap);
log.info("evaluateResult:" + evaluateResult);
// 提取预测结果
String predictedValue = null;
TargetField targetFieldName = evaluator.getTargetField();
Object targetFieldValue = evaluateResult.get(targetFieldName.getFieldName());
// 输出预测结果
if (targetFieldValue instanceof ProbabilityDistribution) {
predictedValue = ((ProbabilityDistribution) targetFieldValue).getPrediction().toString(); log.info("Predicted value(ProbabilityDistribution) : " + predictedValue); } else if (targetFieldValue instanceof FieldValue) { FieldValue fieldValue = (FieldValue) targetFieldValue; predictedValue = fieldValue.asString(); log.info("Predicted value(FieldValue) : " + predictedValue); } else if (targetFieldValue instanceof List) { List resultList = ((List) targetFieldValue)
.stream()
.map(e -> ((FieldValue) e).asString())
.collect(Collectors.toList());
predictedValue = String.join(",", resultList);
log.info("Predicted value(List) : " + predictedValue);
} else {
log.error("unknown type for targetFieldValue:" + targetFieldValue);
}
return predictedValue;
}
}