[lmi] add generated token count to details (#1600)

deepjavalibrary · Mar 5, 2024 · ce4afea · ce4afea
1 parent 350068b
commit ce4afea
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 2 deletions.
diff --git a/engines/python/setup/djl_python/rolling_batch/rolling_batch.py b/engines/python/setup/djl_python/rolling_batch/rolling_batch.py
@@ -88,7 +88,8 @@ def _jsonlines_output_formatter(token: Token, first_token: bool,
         final_dict["generated_text"] = generated_tokens
         if details:
             final_dict["details"] = {
-                "finish_reason": details.get("finish_reason", None)
+                "finish_reason": details.get("finish_reason", None),
+                "generated_tokens": details.get("generated_tokens", None),
             }
     json_encoded_str = json.dumps(final_dict, ensure_ascii=False) + "\n"
     return json_encoded_str
@@ -160,6 +161,7 @@ def set_next_token(self,
             if self.token_cache is not None:
                 details["finish_reason"] = finish_reason
                 details["tokens"] = self.token_cache
+                details["generated_tokens"] = len(self.token_cache)
         if output_formatter is None:
             self.next_token_str = next_token.text
         else:  # output only supports size one now

diff --git a/engines/python/setup/djl_python/tests/test_rolling_batch.py b/engines/python/setup/djl_python/tests/test_rolling_batch.py
@@ -75,6 +75,8 @@ def test_details(self):
             "details": {
                 "finish_reason":
                 "length",
+                "generated_tokens":
+                3,
                 "tokens": [{
                     "id": [244],
                     "text": "He",
@@ -110,7 +112,8 @@ def test_details(self):
             },
             "generated_text": "Hello world",
             "details": {
-                "finish_reason": "length"
+                "finish_reason": "length",
+                "generated_tokens": 3,
             }
         }