[xgb] Add .xgb file extension support (#2810)

deepjavalibrary · Apr 26, 2024 · 7f109c3 · 7f109c3
1 parent f157d01
commit 7f109c3
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbModel.java
@@ -80,6 +80,8 @@ private Path findModelFile(String prefix) {
             String fileName = file.toFile().getName();
             if (fileName.endsWith(".json")) {
                 modelName = fileName.substring(0, fileName.length() - 5);
+            } else if (fileName.endsWith(".xgb")) {
+                modelName = fileName.substring(0, fileName.length() - 4);
             } else {
                 modelName = fileName;
             }
@@ -90,13 +92,22 @@ private Path findModelFile(String prefix) {
         }
         Path modelFile = modelDir.resolve(prefix);
         if (Files.notExists(modelFile) || !Files.isRegularFile(modelFile)) {
-            if (prefix.endsWith(".json")) {
+            if (prefix.endsWith(".json") || prefix.endsWith(".xgb")) {
                 return null;
             }
             modelFile = modelDir.resolve(prefix + ".json");
-            if (Files.notExists(modelFile) || !Files.isRegularFile(modelFile)) {
-                return null;
+            if (Files.isRegularFile(modelFile)) {
+                return modelFile;
+            }
+            modelFile = modelDir.resolve(prefix + ".xgb");
+            if (Files.isRegularFile(modelFile)) {
+                return modelFile;
+            }
+            modelFile = modelDir.resolve("model.xgb");
+            if (Files.isRegularFile(modelFile)) {
+                return modelFile;
             }
+            return null;
         }
         return modelFile;
     }