Add number of tokens to offline (#1623)

mlcommons · Feb 12, 2024 · f9a643c · f9a643c
1 parent f06b920
commit f9a643c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/language/llama2-70b/SUT.py b/language/llama2-70b/SUT.py
@@ -207,9 +207,10 @@ def process_queries(self):
                                                                 query_id_list=query_ids)
 
             for i in range(len(qitem)):
+                n_tokens = processed_output[i].shape[0]
                 response_array = array.array("B", processed_output[i].tobytes())
                 bi = response_array.buffer_info()
-                response = [lg.QuerySampleResponse(qitem[i].id, bi[0], bi[1])]
+                response = [lg.QuerySampleResponse(qitem[i].id, bi[0], bi[1], n_tokens)]
                 lg.QuerySamplesComplete(response)
 
             tok = time.time()