[fasttext] keep __label__ prefix (#1196)

Change-Id: I06ea826ea37d367ea75ca257bd640583b1660fc2
deepjavalibrary · Aug 30, 2021 · 887bc6f · 887bc6f
1 parent e278a05
commit 887bc6f
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 6 deletions.
diff --git a/extensions/fasttext/src/main/java/ai/djl/fasttext/jni/FtWrapper.java b/extensions/fasttext/src/main/java/ai/djl/fasttext/jni/FtWrapper.java
@@ -67,11 +67,7 @@ public Classifications predictProba(String text, int topK) {
         List<String> classes = new ArrayList<>(size);
         List<Double> probabilities = new ArrayList<>(size);
         for (int i = 0; i < size; ++i) {
-            String label = labels[i];
-            if (label.startsWith("__label__")) {
-                label = label.substring(9);
-            }
-            classes.add(label);
+            classes.add(labels[i]); // NOPMD
             probabilities.add((double) probs[i]);
         }
         return new Classifications(classes, probabilities);

diff --git a/extensions/fasttext/src/test/java/ai/djl/fasttext/CookingStackExchangeTest.java b/extensions/fasttext/src/test/java/ai/djl/fasttext/CookingStackExchangeTest.java
@@ -78,7 +78,7 @@ public void testTextClassification()
             String input = "Which baking dish is best to bake a banana bread ?";
             FtModel ftModel = (FtModel) model.getWrappedModel();
             Classifications result = ftModel.classify(input, 8);
-            Assert.assertEquals(result.item(0).getClassName(), "bread");
+            Assert.assertEquals(result.item(0).getClassName(), "__label__bread");
         }
     }