fix

deepjavalibrary · Jul 31, 2023 · 2fb20f0 · 2fb20f0
1 parent 7a4e63f
commit 2fb20f0
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
@@ -340,6 +340,14 @@ public NDArray contrastiveSearch(NDArray inputIds) throws TranslateException {
             }
 
             // TODO: <EOS>, delete the sentence and add it to result.
+            long[] outputIdsArray = searchState.getPastOutputIds().toLongArray();
+            for (int i = 0; i < outputIdsArray.length; i++) {
+                if (outputIdsArray[i] == config.getEosTokenId()) {
+                    if (!exitIndexEndPosition.containsKey((long) i)) {
+                        exitIndexEndPosition.put((long) i, seqLength);
+                    }
+                }
+            }
             if (searchState.getPastOutputIds().getShape().get(1) >= config.getMaxSeqLength()) {
                 break;
             }