From c5662663ceb589b7a44efef97a5ffd6b34244ed1 Mon Sep 17 00:00:00 2001
From: Jorg Doku <jd728@cornell.edu>
Date: Thu, 31 Aug 2023 16:00:28 -0500
Subject: [PATCH] update response

---
 src/handler.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/src/handler.py b/src/handler.py
index 2699a3a..652545c 100644
--- a/src/handler.py
+++ b/src/handler.py
@@ -217,7 +217,15 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
 
         # Record job input and token counts
         runpod_metrics['job_input'] = job_input
-        runpod_metrics['input_tokens'] = len(request_output.prompt_token_ids) * num_seqs
+
+        # Only include the input_tokens count for the very first stream response. This is to avoid duplicate counting.
+        if tracker.stream_index == 0:
+            input_tokens_count = len(request_output.prompt_token_ids)
+            runpod_metrics['input_tokens'] = [input_tokens_count] * num_seqs
+        else:
+            runpod_metrics['input_tokens'] = [0] * num_seqs
+
+        # Include the output tokens count [#, #, #, ...]
         runpod_metrics['output_tokens'] = output_tokens
 
         # Store the scenario type and stream index
@@ -229,6 +237,8 @@ async def handler_streaming(job: dict) -> Generator[dict[str, list], None, None]
 
         ret = {
             "text": text_outputs,
+            "input_tokens": runpod_metrics['input_tokens'],
+            "output_tokens": runpod_metrics['output_tokens'],
             "metrics": runpod_metrics,
         }
         yield ret
@@ -289,6 +299,8 @@ async def handler(job: dict) -> dict[str, list]:
 
     ret = {
         "text": text_outputs,
+        "input_tokens": runpod_metrics['input_tokens'],
+        "output_tokens": runpod_metrics['output_tokens'],
         "metrics": runpod_metrics
     }
     return ret