Repository: deepjavalibrary/djl
Pull request #1414 (2021-12-13 21:17:22)

Create IndexEvaluator and IndexLoss
Right now, there is inconsistent handling for how evaluators and losses can be
applied to only a single array within an NDList. This can happen often when
using multiple losses such as with the AbstractCompositeLoss cases.

Before, we had indexing through SimpleCompositeLoss. However, this would not
work with the trainer evaluators as well. So, the solution must be to move the
tracking of indices into the evaluators/losses.

This creates the indexed variations as a wrapper class so it requires less
updating for existing classes and keeps them simpler. As part of it, the classes
that already had indexing support including some accuracy evaluators and
SimpleCompositeLoss had it removed.


Patch:
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
index 1f9d3db0a..4a2e34a10 100644
--- a/.github/CODEOWNERS
+++ b/.github/CODEOWNERS
@@ -1,2 +1,2 @@
 # Global Owners
-* @zachgk @frankfliu
\ No newline at end of file
+* @zachgk @frankfliu @DeepJavaLibrary/djl-admin
\ No newline at end of file
diff --git a/.github/workflows/continuous.yml b/.github/workflows/continuous.yml
index 1f4b381b3..71d79e73e 100644
--- a/.github/workflows/continuous.yml
+++ b/.github/workflows/continuous.yml
@@ -73,6 +73,8 @@ jobs:
           key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
           restore-keys: |
             ${{ runner.os }}-gradle-
+      - name: Compile Sentencepiece JNI
+        run: ./gradlew :extensions:sentencepiece:compileJNI
       - name: Compile PyTorch JNI
         run: ./gradlew :engines:pytorch:pytorch-native:compileJNI
       - name: Compile PaddlePaddle JNI
diff --git a/.github/workflows/native_s3_sentencepiece.yml b/.github/workflows/native_s3_sentencepiece.yml
index 27557ca24..962b3ad58 100644
--- a/.github/workflows/native_s3_sentencepiece.yml
+++ b/.github/workflows/native_s3_sentencepiece.yml
@@ -71,3 +71,34 @@ jobs:
           SENTENCEPIECE_VERSION="$(cat gradle.properties | awk -F '=' '/sentencepiece_version/ {print $2}')"
           aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/*"
+
+  build-sentencepiece-jni-windows:
+    runs-on: windows-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up JDK 1.8
+        uses: actions/setup-java@v1
+        with:
+          java-version: 1.8
+      - uses: actions/cache@v2
+        with:
+          path: ~/.gradle/caches
+          key: ${{ runner.os }}-gradle-${{ hashFiles('**/*.gradle*') }}
+          restore-keys: |
+            ${{ runner.os }}-gradle-
+      - name: Release JNI prep
+        run: |
+          ./gradlew :extensions:sentencepiece:compileJNI
+          ./gradlew -Pjni :extensions:sentencepiece:test
+      - name: Configure AWS Credentials
+        uses: aws-actions/configure-aws-credentials@v1
+        with:
+          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
+          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
+          aws-region: us-east-2
+      - name: Copy files to S3 with the AWS CLI
+        shell: bash
+        run: |
+          SENTENCEPIECE_VERSION="$(cat gradle.properties | awk -F '=' '/sentencepiece_version/ {print $2}')"
+          aws s3 sync extensions/sentencepiece/jnilib s3://djl-ai/publish/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/sentencepiece-${SENTENCEPIECE_VERSION}/jnilib/*"
diff --git a/api/src/main/java/ai/djl/training/DefaultTrainingConfig.java b/api/src/main/java/ai/djl/training/DefaultTrainingConfig.java
index 35b25683a..1218404e9 100644
--- a/api/src/main/java/ai/djl/training/DefaultTrainingConfig.java
+++ b/api/src/main/java/ai/djl/training/DefaultTrainingConfig.java
@@ -24,6 +24,7 @@ import ai.djl.training.optimizer.Optimizer;
 import ai.djl.util.PairList;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.Collection;
 import java.util.List;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ForkJoinPool;
@@ -136,6 +137,18 @@ public class DefaultTrainingConfig implements TrainingConfig {
         return this;
     }
 
+    /**
+     * Adds multiple {@link Evaluator}s that needs to be computed during training.
+     *
+     * @param evaluators the evaluators to be added
+     * @param <T> the type of evaluator to be added
+     * @return this {@code DefaultTrainingConfig}
+     */
+    public <T extends Evaluator> DefaultTrainingConfig addEvaluators(Collection<T> evaluators) {
+        evaluators.forEach(this::addEvaluator);
+        return this;
+    }
+
     /**
      * Adds an {@link Evaluator} that needs to be computed during training.
      *
diff --git a/api/src/main/java/ai/djl/training/evaluator/AbstractAccuracy.java b/api/src/main/java/ai/djl/training/evaluator/AbstractAccuracy.java
index acfba1fff..b58dcb9c2 100644
--- a/api/src/main/java/ai/djl/training/evaluator/AbstractAccuracy.java
+++ b/api/src/main/java/ai/djl/training/evaluator/AbstractAccuracy.java
@@ -29,30 +29,26 @@ public abstract class AbstractAccuracy extends Evaluator {
 
     protected Map<String, Long> correctInstances;
     protected int axis;
-    protected int index;
 
     /**
-     * Creates an accuracy evaluator that computes accuracy across axis 1 along given index.
+     * Creates an accuracy evaluator that computes accuracy across axis 1.
      *
      * @param name the name of the evaluator, default is "Accuracy"
-     * @param index the index of the NDArray in labels to compute accuracy for
      */
-    public AbstractAccuracy(String name, int index) {
-        this(name, index, 1);
+    public AbstractAccuracy(String name) {
+        this(name, 1);
     }
 
     /**
      * Creates an accuracy evaluator.
      *
      * @param name the name of the evaluator, default is "Accuracy"
-     * @param index the index of the NDArray in labels to compute accuracy for
      * @param axis the axis that represent classes in prediction, default 1
      */
-    public AbstractAccuracy(String name, int index, int axis) {
+    public AbstractAccuracy(String name, int axis) {
         super(name);
         correctInstances = new ConcurrentHashMap<>();
         this.axis = axis;
-        this.index = index;
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/training/evaluator/Accuracy.java b/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
index 5c4fe5ed6..5b9695243 100644
--- a/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
+++ b/api/src/main/java/ai/djl/training/evaluator/Accuracy.java
@@ -25,36 +25,34 @@ public class Accuracy extends AbstractAccuracy {
      * index.
      */
     public Accuracy() {
-        this("Accuracy", 0, 1);
+        this("Accuracy", 1);
     }
 
     /**
-     * Creates a multiclass accuracy evaluator that computes accuracy across axis 1 along given
+     * Creates a multiclass accuracy evaluator that computes accuracy across axis 1 along the 0th
      * index.
      *
      * @param name the name of the evaluator, default is "Accuracy"
-     * @param index the index of the NDArray in labels to compute accuracy for
      */
-    public Accuracy(String name, int index) {
-        super(name, index);
+    public Accuracy(String name) {
+        this(name, 1);
     }
 
     /**
      * Creates a multiclass accuracy evaluator.
      *
      * @param name the name of the evaluator, default is "Accuracy"
-     * @param index the index of the NDArray in labels to compute accuracy for
      * @param axis the axis that represent classes in prediction, default 1
      */
-    public Accuracy(String name, int index, int axis) {
-        super(name, index, axis);
+    public Accuracy(String name, int axis) {
+        super(name, axis);
     }
 
     /** {@inheritDoc} */
     @Override
     protected Pair<Long, NDArray> accuracyHelper(NDList labels, NDList predictions) {
-        NDArray label = labels.get(index);
-        NDArray prediction = predictions.get(index);
+        NDArray label = labels.head();
+        NDArray prediction = predictions.head();
         checkLabelShapes(label, prediction);
         NDArray predictionReduced;
         if (!label.getShape().equals(prediction.getShape())) {
diff --git a/api/src/main/java/ai/djl/training/evaluator/BinaryAccuracy.java b/api/src/main/java/ai/djl/training/evaluator/BinaryAccuracy.java
index 42bf45b0f..3468be2d4 100644
--- a/api/src/main/java/ai/djl/training/evaluator/BinaryAccuracy.java
+++ b/api/src/main/java/ai/djl/training/evaluator/BinaryAccuracy.java
@@ -35,11 +35,10 @@ public class BinaryAccuracy extends AbstractAccuracy {
      *
      * @param name the name of the evaluator, default is "Accuracy"
      * @param threshold the value differentiating the posive and negative classes (usually 0 or .5)
-     * @param index the index of the NDArray in labels to compute accuracy for
      * @param axis the axis that represent classes in prediction, default 1
      */
-    public BinaryAccuracy(String name, float threshold, int index, int axis) {
-        super(name, index, axis);
+    public BinaryAccuracy(String name, float threshold, int axis) {
+        super(name, axis);
         this.threshold = threshold;
     }
 
@@ -49,10 +48,9 @@ public class BinaryAccuracy extends AbstractAccuracy {
      *
      * @param name the name of the evaluator, default is "Accuracy"
      * @param threshold the value differentiating the posive and negative classes (usually 0 or .5)
-     * @param index the index of the NDArray in labels to compute accuracy for
      */
-    public BinaryAccuracy(String name, float threshold, int index) {
-        this(name, threshold, index, 1);
+    public BinaryAccuracy(String name, float threshold) {
+        this(name, threshold, 1);
     }
 
     /**
@@ -62,7 +60,7 @@ public class BinaryAccuracy extends AbstractAccuracy {
      * @param threshold the value differentiating the posive and negative classes (usually 0 or .5)
      */
     public BinaryAccuracy(float threshold) {
-        this("BinaryAccuracy", threshold, 0, 1);
+        this("BinaryAccuracy", threshold, 1);
     }
 
     /** Creates a binary (two class) accuracy evaluator with 0 threshold. */
@@ -76,8 +74,8 @@ public class BinaryAccuracy extends AbstractAccuracy {
         Preconditions.checkArgument(
                 labels.size() == predictions.size(),
                 "labels and prediction length does not match.");
-        NDArray label = labels.get(index);
-        NDArray prediction = predictions.get(index);
+        NDArray label = labels.head();
+        NDArray prediction = predictions.head();
         checkLabelShapes(label, prediction, false);
         NDArray predictionReduced = prediction.gte(threshold);
         // result of sum is int64 now
diff --git a/api/src/main/java/ai/djl/training/evaluator/IndexEvaluator.java b/api/src/main/java/ai/djl/training/evaluator/IndexEvaluator.java
new file mode 100644
index 000000000..a7fe08b61
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/evaluator/IndexEvaluator.java
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.evaluator;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * A wrapper for an {@link Evaluator} that evaluates on only a particular {@link NDArray} in the
+ * predictions and/or labels {@link NDList}s.
+ */
+public class IndexEvaluator extends Evaluator {
+
+    private Evaluator evaluator;
+    private Integer predictionsIndex;
+    private Integer labelsIndex;
+
+    /**
+     * Constructs an {@link IndexEvaluator} with the same index for both predictions and labels.
+     *
+     * @param evaluator the base evaluator
+     * @param index the index for both predictions and labels
+     */
+    public IndexEvaluator(Evaluator evaluator, int index) {
+        this(evaluator, index, index);
+    }
+
+    /**
+     * Constructs an {@link IndexEvaluator}.
+     *
+     * @param evaluator the base evaluator
+     * @param predictionsIndex the predictions index
+     * @param labelsIndex the labels index
+     */
+    public IndexEvaluator(Evaluator evaluator, Integer predictionsIndex, Integer labelsIndex) {
+        super(evaluator.getName());
+        this.evaluator = evaluator;
+        this.predictionsIndex = predictionsIndex;
+        this.labelsIndex = labelsIndex;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        return evaluator.evaluate(getLabels(labels), getPredictions(predictions));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void addAccumulator(String key) {
+        evaluator.addAccumulator(key);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void updateAccumulator(String key, NDList labels, NDList predictions) {
+        evaluator.updateAccumulator(key, getLabels(labels), getPredictions(predictions));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void resetAccumulator(String key) {
+        evaluator.resetAccumulator(key);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public float getAccumulator(String key) {
+        return evaluator.getAccumulator(key);
+    }
+
+    private NDList getPredictions(NDList predictions) {
+        if (predictionsIndex == null) {
+            return predictions;
+        }
+        return new NDList(predictions.get(predictionsIndex));
+    }
+
+    private NDList getLabels(NDList labels) {
+        if (labelsIndex == null) {
+            return labels;
+        }
+        return new NDList(labels.get(labelsIndex));
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/evaluator/TopKAccuracy.java b/api/src/main/java/ai/djl/training/evaluator/TopKAccuracy.java
index f7b56acf1..4a4fa0210 100644
--- a/api/src/main/java/ai/djl/training/evaluator/TopKAccuracy.java
+++ b/api/src/main/java/ai/djl/training/evaluator/TopKAccuracy.java
@@ -34,11 +34,10 @@ public class TopKAccuracy extends AbstractAccuracy {
      * Creates a {@code TopKAccuracy} instance.
      *
      * @param name the accuracy name, default "Top_K_Accuracy"
-     * @param index the index of the {@link NDArray} in labels to compute topK accuracy for
      * @param topK the value of K
      */
-    public TopKAccuracy(String name, int index, int topK) {
-        super(name, index);
+    public TopKAccuracy(String name, int topK) {
+        super(name);
         if (topK > 1) {
             this.topK = topK;
         } else {
@@ -46,17 +45,6 @@ public class TopKAccuracy extends AbstractAccuracy {
         }
     }
 
-    /**
-     * Creates an instance of {@code TopKAccuracy} evaluator that computes topK accuracy across axis
-     * 1 along the given index.
-     *
-     * @param index the index of the {@link NDArray} in labels to compute topK accuracy for
-     * @param topK the value of K
-     */
-    public TopKAccuracy(int index, int topK) {
-        this("Top_" + topK + "_Accuracy", index, topK);
-    }
-
     /**
      * Creates an instance of {@code TopKAccuracy} evaluator that computes topK accuracy across axis
      * 1 along the 0th index.
@@ -64,14 +52,14 @@ public class TopKAccuracy extends AbstractAccuracy {
      * @param topK the value of K
      */
     public TopKAccuracy(int topK) {
-        this("Top_" + topK + "_Accuracy", 0, topK);
+        this("Top_" + topK + "_Accuracy", topK);
     }
 
     /** {@inheritDoc} */
     @Override
     protected Pair<Long, NDArray> accuracyHelper(NDList labels, NDList predictions) {
-        NDArray label = labels.get(index);
-        NDArray prediction = predictions.get(index);
+        NDArray label = labels.head();
+        NDArray prediction = predictions.head();
         // number of labels and predictions should be the same
         checkLabelShapes(label, prediction);
         // ascending by default
diff --git a/api/src/main/java/ai/djl/training/loss/IndexLoss.java b/api/src/main/java/ai/djl/training/loss/IndexLoss.java
new file mode 100644
index 000000000..4f05c8781
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/loss/IndexLoss.java
@@ -0,0 +1,71 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.loss;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * A wrapper for a {@link Loss} that evaluates on only a particular {@link NDArray} in the
+ * predictions and/or labels {@link NDList}s.
+ */
+public class IndexLoss extends Loss {
+
+    private Loss loss;
+    private Integer predictionsIndex;
+    private Integer labelsIndex;
+
+    /**
+     * Constructs an {@link IndexLoss} with the same index for both predictions and labels.
+     *
+     * @param loss the base evaluator
+     * @param index the index for both predictions and labels
+     */
+    public IndexLoss(Loss loss, int index) {
+        this(loss, index, index);
+    }
+
+    /**
+     * Constructs an {@link IndexLoss}.
+     *
+     * @param loss the base evaluator
+     * @param predictionsIndex the predictions index
+     * @param labelsIndex the labels index
+     */
+    public IndexLoss(Loss loss, Integer predictionsIndex, Integer labelsIndex) {
+        super(loss.getName());
+        this.loss = loss;
+        this.predictionsIndex = predictionsIndex;
+        this.labelsIndex = labelsIndex;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        return loss.evaluate(getLabels(labels), getPredictions(predictions));
+    }
+
+    private NDList getPredictions(NDList predictions) {
+        if (predictionsIndex == null) {
+            return predictions;
+        }
+        return new NDList(predictions.get(predictionsIndex));
+    }
+
+    private NDList getLabels(NDList labels) {
+        if (labelsIndex == null) {
+            return labels;
+        }
+        return new NDList(labels.get(labelsIndex));
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/loss/SimpleCompositeLoss.java b/api/src/main/java/ai/djl/training/loss/SimpleCompositeLoss.java
index ccaa46d7c..2dc0519aa 100644
--- a/api/src/main/java/ai/djl/training/loss/SimpleCompositeLoss.java
+++ b/api/src/main/java/ai/djl/training/loss/SimpleCompositeLoss.java
@@ -15,15 +15,13 @@ package ai.djl.training.loss;
 import ai.djl.ndarray.NDList;
 import ai.djl.util.Pair;
 import java.util.ArrayList;
-import java.util.List;
 
 /**
  * {@code SimpleCompositeLoss} is an implementation of the {@link Loss} abstract class that can
  * combine different {@link Loss} functions by adding the individual losses together.
  *
- * <p>This class can be used when the losses either accept a single index of the labels and
- * predictions or the entire lists. For more complicated composite losses, extend the {@link
- * AbstractCompositeLoss}.
+ * <p>For cases where the losses use only a single index of the labels and/or predictions, use the
+ * {@link IndexLoss}.
  *
  * <p>For an example of using this loss, see <a
  * href="https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/training/TrainCaptcha.java">the
@@ -31,8 +29,6 @@ import java.util.List;
  */
 public class SimpleCompositeLoss extends AbstractCompositeLoss {
 
-    private List<Integer> indices;
-
     /**
      * Creates a new empty instance of {@code CompositeLoss} that can combine the given {@link Loss}
      * components.
@@ -50,7 +46,6 @@ public class SimpleCompositeLoss extends AbstractCompositeLoss {
     public SimpleCompositeLoss(String name) {
         super(name);
         components = new ArrayList<>();
-        indices = new ArrayList<>();
     }
 
     /**
@@ -61,21 +56,6 @@ public class SimpleCompositeLoss extends AbstractCompositeLoss {
      */
     public SimpleCompositeLoss addLoss(Loss loss) {
         components.add(loss);
-        indices.add(null);
-        return this;
-    }
-
-    /**
-     * Adds a Loss that applies to a single index of the label and predictions to this composite
-     * loss.
-     *
-     * @param loss the loss to add
-     * @param index the index in the label and predictions NDLists this loss applies to
-     * @return this composite loss
-     */
-    public SimpleCompositeLoss addLoss(Loss loss, int index) {
-        components.add(loss);
-        indices.add(index);
         return this;
     }
 
@@ -83,11 +63,6 @@ public class SimpleCompositeLoss extends AbstractCompositeLoss {
     @Override
     protected Pair<NDList, NDList> inputForComponent(
             int componentIndex, NDList labels, NDList predictions) {
-        if (indices.get(componentIndex) != null) {
-            int index = indices.get(componentIndex);
-            return new Pair<>(new NDList(labels.get(index)), new NDList(predictions.get(index)));
-        } else {
-            return new Pair<>(labels, predictions);
-        }
+        return new Pair<>(labels, predictions);
     }
 }
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainCaptcha.java b/examples/src/main/java/ai/djl/examples/training/TrainCaptcha.java
index 766e16192..184ff3908 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainCaptcha.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainCaptcha.java
@@ -31,8 +31,10 @@ import ai.djl.training.dataset.Dataset;
 import ai.djl.training.dataset.Dataset.Usage;
 import ai.djl.training.dataset.RandomAccessDataset;
 import ai.djl.training.evaluator.Accuracy;
+import ai.djl.training.evaluator.IndexEvaluator;
 import ai.djl.training.listener.SaveModelTrainingListener;
 import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.IndexLoss;
 import ai.djl.training.loss.SimpleCompositeLoss;
 import ai.djl.training.loss.SoftmaxCrossEntropyLoss;
 import ai.djl.training.util.ProgressBar;
@@ -99,17 +101,18 @@ public final class TrainCaptcha {
                 });
         SimpleCompositeLoss loss = new SimpleCompositeLoss();
         for (int i = 0; i < CaptchaDataset.CAPTCHA_LENGTH; i++) {
-            loss.addLoss(new SoftmaxCrossEntropyLoss("loss_digit_" + i), i);
+            loss.addLoss(new IndexLoss(new SoftmaxCrossEntropyLoss("loss_digit_" + i), i));
         }
 
         DefaultTrainingConfig config =
                 new DefaultTrainingConfig(loss)
                         .optDevices(Engine.getInstance().getDevices(arguments.getMaxGpus()))
+                        .addEvaluators(loss.getComponents())
                         .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
                         .addTrainingListeners(listener);
 
         for (int i = 0; i < CaptchaDataset.CAPTCHA_LENGTH; i++) {
-            config.addEvaluator(new Accuracy("acc_digit_" + i, i));
+            config.addEvaluator(new IndexEvaluator(new Accuracy("acc_digit_" + i), i));
         }
 
         return config;
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainSeq2Seq.java b/examples/src/main/java/ai/djl/examples/training/TrainSeq2Seq.java
index 6fd7e2c33..2e0804dcc 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainSeq2Seq.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainSeq2Seq.java
@@ -149,7 +149,7 @@ public final class TrainSeq2Seq {
                 });
 
         return new DefaultTrainingConfig(new MaskedSoftmaxCrossEntropyLoss())
-                .addEvaluator(new Accuracy("Accuracy", 0, 2))
+                .addEvaluator(new Accuracy("Accuracy", 2))
                 .optDevices(Engine.getInstance().getDevices(arguments.getMaxGpus()))
                 .optExecutorService(executorService)
                 .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
diff --git a/extensions/sentencepiece/CMakeLists.txt b/extensions/sentencepiece/CMakeLists.txt
index 54566a936..7aea0fd21 100644
--- a/extensions/sentencepiece/CMakeLists.txt
+++ b/extensions/sentencepiece/CMakeLists.txt
@@ -9,6 +9,14 @@ set(SPM_ENABLE_TCMALLOC OFF CACHE BOOL "Build sentencepiece static library")
 # disable shared library as we only link static library
 set(SPM_ENABLE_SHARED OFF CACHE BOOL "Build sentencepiece static library")
 
+if (MSVC)
+  string(REPLACE "/MD" "/MT" CMAKE_CXX_FLAGS_DEBUG          ${CMAKE_CXX_FLAGS_DEBUG})
+  string(REPLACE "/MD" "/MT" CMAKE_CXX_FLAGS_MINSIZEREL     ${CMAKE_CXX_FLAGS_MINSIZEREL})
+  string(REPLACE "/MD" "/MT" CMAKE_CXX_FLAGS_RELEASE        ${CMAKE_CXX_FLAGS_RELEASE})
+  string(REPLACE "/MD" "/MT" CMAKE_CXX_FLAGS_RELWITHDEBINFO ${CMAKE_CXX_FLAGS_RELWITHDEBINFO})
+  add_definitions("/wd4267 /wd4244 /wd4305 /Zc:strictStrings /utf-8")
+endif()
+
 find_package(JNI REQUIRED)
 find_path(UTILS_INCLUDE_DIR NAMES djl/utils.h PATHS ${PROJECT_SOURCE_DIR}/../../api/src/main/native REQUIRED)
 
diff --git a/extensions/sentencepiece/build.cmd b/extensions/sentencepiece/build.cmd
new file mode 100644
index 000000000..1253de17c
--- /dev/null
+++ b/extensions/sentencepiece/build.cmd
@@ -0,0 +1,25 @@
+set VERSION="%1"
+set FILEPATH="sentencepiece"
+
+if exist %FILEPATH% (
+    echo Found %FILEPATH%
+) else (
+    echo Couldn't find %FILEPATH%, downloading it ...
+    echo Cloning google/sentencepiece.git
+    git clone https://github.com/google/sentencepiece.git -b %VERSION%
+    echo Finished cloning sentencepiece
+)
+
+if exist build rd /q /s build
+md build\classes
+cd build
+
+javac -sourcepath ../src/main/java/ ../src/main/java/ai/djl/sentencepiece/jni/SentencePieceLibrary.java -h include -d classes
+cmake ..
+cmake --build . --config Release
+
+:: for nightly ci
+md jnilib\win-x86_64
+copy Release\sentencepiece_native.dll jnilib\win-x86_64
+
+cd ..
diff --git a/extensions/sentencepiece/build.gradle b/extensions/sentencepiece/build.gradle
index e46229873..7dd054cff 100644
--- a/extensions/sentencepiece/build.gradle
+++ b/extensions/sentencepiece/build.gradle
@@ -19,6 +19,7 @@ processResources {
     doLast {
         def url = "https://publish.djl.ai/sentencepiece-${sentencepiece_version}/jnilib/${djl_version}"
         def files = [
+                "win-x86_64"  : "sentencepiece_native.dll",
                 "linux-x86_64": "libsentencepiece_native.so",
                 "osx-x86_64"  : "libsentencepiece_native.dylib"
         ]
@@ -61,22 +62,26 @@ apply from: file("${rootProject.projectDir}/tools/gradle/cpp-formatter.gradle")
 
 task compileJNI {
     doFirst {
-        if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("mac")
-                || System.properties['os.name'].toLowerCase(Locale.ROOT).contains("linux")) {
+        if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("win")) {
             exec {
-                commandLine 'bash', 'build.sh'
+                commandLine "${project.projectDir}/build.cmd", "v${sentencepiece_version}"
             }
-
-            // for ci to upload to S3
-            def ciDir = "${project.projectDir}/jnilib/${djl_version}/"
-            copy {
-                from "${project.buildDir}/jnilib"
-                into ciDir
+        } else if (System.properties['os.name'].toLowerCase(Locale.ROOT).contains("mac")
+                || System.properties['os.name'].toLowerCase(Locale.ROOT).contains("linux")) {
+            exec {
+                commandLine "bash", "build.sh", "v${sentencepiece_version}"
             }
-            delete System.getProperty("user.home") + "/.djl.ai/sentencepiece"
         } else {
             throw new IllegalStateException("Unknown Architecture " + System.properties['os.name'])
         }
+
+        // for ci to upload to S3
+        def ciDir = "${project.projectDir}/jnilib/${djl_version}/"
+        copy {
+            from "${project.buildDir}/jnilib"
+            into ciDir
+        }
+        delete System.getProperty("user.home") + "/.djl.ai/sentencepiece"
     }
 }
 
diff --git a/extensions/sentencepiece/build.sh b/extensions/sentencepiece/build.sh
index e82dd9198..c1ddb9852 100755
--- a/extensions/sentencepiece/build.sh
+++ b/extensions/sentencepiece/build.sh
@@ -10,7 +10,7 @@ elif [[ -n $(command -v sysctl) ]]; then
 fi
 PLATFORM=$(uname | tr '[:upper:]' '[:lower:]')
 
-VERSION=v0.1.95
+VERSION=$1
 
 pushd $WORK_DIR
 if [ ! -d "sentencepiece" ];
diff --git a/extensions/sentencepiece/src/main/java/ai/djl/sentencepiece/jni/LibUtils.java b/extensions/sentencepiece/src/main/java/ai/djl/sentencepiece/jni/LibUtils.java
index 7a51d7471..09821b9cd 100644
--- a/extensions/sentencepiece/src/main/java/ai/djl/sentencepiece/jni/LibUtils.java
+++ b/extensions/sentencepiece/src/main/java/ai/djl/sentencepiece/jni/LibUtils.java
@@ -43,10 +43,6 @@ public final class LibUtils {
     private LibUtils() {}
 
     public static void loadLibrary() {
-        if (System.getProperty("os.name").startsWith("Win")) {
-            throw new UnsupportedOperationException("Windows is not supported.");
-        }
-
         String libName = copyJniLibraryFromClasspath();
         logger.debug("Loading sentencepiece library from: {}", libName);
         System.load(libName); // NOPMD
diff --git a/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc b/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
index cf6d1c73c..de0d94261 100644
--- a/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
+++ b/extensions/sentencepiece/src/main/native/ai_djl_sentencepiece_jni_SentencePieceLibrary.cc
@@ -96,7 +96,7 @@ JNIEXPORT jstring JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_idT
   return env->NewStringUTF(processor_ptr->IdToPiece(jid).c_str());
 }
 
-JNIEXPORT int JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_pieceToId(
+JNIEXPORT jint JNICALL Java_ai_djl_sentencepiece_jni_SentencePieceLibrary_pieceToId(
     JNIEnv* env, jobject jthis, jlong jhandle, jstring jpiece) {
   auto* processor_ptr = reinterpret_cast<sentencepiece::SentencePieceProcessor*>(jhandle);
   return processor_ptr->PieceToId(djl::utils::jni::GetStringFromJString(env, jpiece));
diff --git a/gradle.properties b/gradle.properties
index 0b1f66f09..b144b7042 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -27,7 +27,7 @@ commons_logging_version=1.2
 gson_version=2.8.8
 jna_version=5.9.0
 slf4j_version=1.7.32
-log4j_slf4j_version=2.15.0
+log4j_slf4j_version=2.17.0
 awssdk_version=2.17.54
 hadoop_version=3.3.1
 javacpp_version=1.5.6


Test patch:
diff --git a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTextEmbeddingTest.java b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTextEmbeddingTest.java
index 811e1e216..18c4fb8c3 100644
--- a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTextEmbeddingTest.java
+++ b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTextEmbeddingTest.java
@@ -12,7 +12,6 @@
  */
 package ai.djl.sentencepiece;
 
-import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.DownloadUtils;
 import java.io.IOException;
 import java.nio.file.Files;
@@ -36,8 +35,6 @@ public class SpTextEmbeddingTest {
 
     @Test
     public void testEmbedding() throws IOException {
-        TestRequirements.notWindows();
-
         Path modelPath = Paths.get("build/test/models/sententpiece_test_model.model");
         try (SpTokenizer tokenizer = new SpTokenizer(modelPath)) {
             SpTextEmbedding embedding = SpTextEmbedding.from(tokenizer);
diff --git a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTokenizerTest.java b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTokenizerTest.java
index e14360390..155290e37 100644
--- a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTokenizerTest.java
+++ b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpTokenizerTest.java
@@ -39,8 +39,6 @@ public class SpTokenizerTest {
 
     @Test
     public void testLoadFromBytes() throws IOException {
-        TestRequirements.notWindows();
-
         Path modelPath = Paths.get("build/test/models/sententpiece_test_model.model");
         byte[] bytes = Files.readAllBytes(modelPath);
         try (SpTokenizer tokenizer = new SpTokenizer(bytes)) {
diff --git a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpVocabularyTest.java b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpVocabularyTest.java
index ec3eaddd6..fbc4d4970 100644
--- a/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpVocabularyTest.java
+++ b/extensions/sentencepiece/src/test/java/ai/djl/sentencepiece/SpVocabularyTest.java
@@ -13,7 +13,6 @@
 
 package ai.djl.sentencepiece;
 
-import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.DownloadUtils;
 import java.io.IOException;
 import java.nio.file.Files;
@@ -37,8 +36,6 @@ public class SpVocabularyTest {
 
     @Test
     public void testTokenIdConversion() throws IOException {
-        TestRequirements.notWindows();
-
         Path modelPath = Paths.get("build/test/models");
         String prefix = "sententpiece_test_model";
         SpTokenizer tokenizer = new SpTokenizer(modelPath, prefix);