[python] Add generated_text key to rolling batch output

deepjavalibrary · Jul 5, 2023 · 173cb55 · 173cb55
1 parent 675c009
commit 173cb55
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java b/engines/python/src/main/java/ai/djl/python/engine/RollingBatch.java
@@ -179,7 +179,10 @@ void addResponse(String json, boolean enableStreaming) {
             } else {
                 nextToken.append(element.get("data").getAsString());
                 if (last) {
-                    data.appendContent(BytesSupplier.wrap(nextToken.toString()), true);
+                    data.appendContent(
+                            BytesSupplier.wrap(
+                                    "{\"generated_text\": [\"" + nextToken.toString() + "\"]}"),
+                            true);
                 }
             }
         }

diff --git a/engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java b/engines/python/src/test/java/ai/djl/python/engine/PyEngineTest.java
@@ -427,7 +427,7 @@ public void testRollingBatch() throws TranslateException, IOException, ModelExce
             Assert.assertNull(cbs.pollChunk());
             String ret = cbs.getAsString();
             System.out.println(ret);
-            Assert.assertTrue(ret.startsWith(" token_request4_"));
+            Assert.assertTrue(ret.startsWith("{\"generated_text\": [\" token_request4_"));
         }
     }