blob: 6a1e2f2562ad8e1a99dc4bb535dcec8c8d2c45c5 (
plain) (
blame)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
package ai.vespa.rankingexpression.importer.lightgbm;
import com.yahoo.json.Jackson;
import ai.vespa.rankingexpression.importer.ImportedModel;
import ai.vespa.rankingexpression.importer.ModelImporter;
import ai.vespa.rankingexpression.importer.configmodelview.ImportedMlModel;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.core.JsonToken;
import com.yahoo.searchlib.rankingexpression.RankingExpression;
import com.yahoo.searchlib.rankingexpression.parser.ParseException;
import java.io.File;
import java.io.IOException;
/**
* Converts a LightGBM model into a ranking expression.
*
* @author lesters
*/
public class LightGBMImporter extends ModelImporter {
@Override
public boolean canImport(String modelPath) {
File modelFile = new File(modelPath);
if ( ! modelFile.isFile()) return false;
return modelFile.toString().endsWith(".json") && probe(modelFile);
}
/**
* Returns true if the give file looks like a LightGBM json file.
* Currently, we just check if the json has an element called "tree_info"
*/
private boolean probe(File modelFile) {
try (JsonParser parser = Jackson.mapper().createParser(modelFile)) {
while (parser.nextToken() != null) {
JsonToken token = parser.currentToken();
if (token == JsonToken.FIELD_NAME) {
if ("tree_info".equals(parser.currentName())) return true;
}
}
return false;
} catch (IOException e) {
throw new IllegalArgumentException("Could not read '" + modelFile + "'", e);
}
}
@Override
public ImportedModel importModel(String modelName, String modelPath) {
try {
ImportedModel model = new ImportedModel(modelName, modelPath, ImportedMlModel.ModelType.LIGHTGBM);
LightGBMParser parser = new LightGBMParser(modelPath);
RankingExpression expression = new RankingExpression(parser.toRankingExpression());
model.expression(modelName, expression);
return model;
} catch (IOException e) {
throw new IllegalArgumentException("Could not import LightGBM model from '" + modelPath + "'", e);
} catch (ParseException e) {
throw new IllegalArgumentException("Could not parse ranking expression resulting from '" + modelPath + "'", e);
}
}
}
|
