[NeuralChat] Support compatible stats format (#1112)

Author: lvliang-intel

intel_extension_for_transformers/neural_chat/config.py

@@ -406,6 +406,8 @@ class GenerationConfig:
  max_gpu_memory: int = None
  use_fp16: bool = False
  ipex_int8: bool = False
+ return_stats: bool = False
+ format_version: str = "v2"
  task: str = ""
 
 @dataclass

intel_extension_for_transformers/neural_chat/models/base_model.py

@@ -52,6 +52,8 @@ def construct_parameters(query, model_name, device, assistant_model, config):
  params["use_hpu_graphs"] = config.use_hpu_graphs
  params["use_cache"] = config.use_cache
  params["ipex_int8"] = config.ipex_int8
+ params["return_stats"] = config.return_stats
+ params["format_version"] = config.format_version
  params["assistant_model"] = assistant_model
  params["device"] = device
  return params

intel_extension_for_transformers/neural_chat/models/model_utils.py

@@ -794,6 +794,7 @@ def predict_stream(**params):
  Determines whether to utilize Habana Processing Units (HPUs) for accelerated generation.
  `use_cache` (bool): Determines whether to utilize kv cache for accelerated generation.
  `ipex_int8` (bool): Whether to use IPEX int8 model to inference.
+ `format_version` (string): the format version of return stats.
 
  Returns:
  generator: A generator that yields the generated streaming text.
@@ -822,6 +823,7 @@ def predict_stream(**params):
  use_hpu_graphs = params["use_hpu_graphs"] if "use_hpu_graphs" in params else False
  use_cache = params["use_cache"] if "use_cache" in params else True
  return_stats = params["return_stats"] if "return_stats" in params else False
+ format_version = params["format_version"] if "format_version" in params else "v2"
  prompt = params["prompt"]
  ipex_int8 = params["ipex_int8"] if "ipex_int8" in params else False
  model = MODELS[model_name]["model"]
@@ -1017,17 +1019,27 @@ def generate_output():
  0
  )
  if return_stats:
- stats = {
- "input_token_len": str(input_token_len),
- "output_token_len": str(output_token_len),
- "duration": str(duration) + " ms",
- "first_token_latency": str(first_token_latency) + " ms",
- "msecond_per_token": str(msecond_per_token) + " ms",
- }
- yield "\n| {:<22} | {:<27} |\n".format("Key", "Value")
- yield "| " + "-"*22 + " | " + "-"*27 + " |" + "\n"
- for key, value in stats.items():
- yield "| {:<22} | {:<27} |\n".format(key, value)
+ if format_version == "v1":
+ stats = {
+ "input_token_len": input_token_len,
+ "output_token_len": output_token_len,
+ "duration": duration,
+ "first_token_latency": first_token_latency,
+ "msecond_per_token": msecond_per_token,
+ }
+ yield "END_OF_STREAM_STATS={}".format(stats)
+ else:
+ stats = {
+ "input_token_len": str(input_token_len),
+ "output_token_len": str(output_token_len),
+ "duration": str(duration) + " ms",
+ "first_token_latency": str(first_token_latency) + " ms",
+ "msecond_per_token": str(msecond_per_token) + " ms",
+ }
+ yield "\n| {:<22} | {:<27} |\n".format("Key", "Value")
+ yield "| " + "-"*22 + " | " + "-"*27 + " |" + "\n"
+ for key, value in stats.items():
+ yield "| {:<22} | {:<27} |\n".format(key, value)
 
 
 def predict(**params):

intel_extension_for_transformers/neural_chat/tests/ci/api/test_chatbot_build_api.py

@@ -167,6 +167,27 @@ def _run_retrieval(local_dir):
  _run_retrieval(local_dir="/tf_dataset2/inc-ut/instructor-large")
  _run_retrieval(local_dir="/tf_dataset2/inc-ut/bge-base-en-v1.5")
 
+ def test_text_chat_stream_return_stats_with_v1_format(self):
+ config = PipelineConfig(model_name_or_path="facebook/opt-125m")
+ chatbot = build_chatbot(config)
+ stream_text = ""
+ gen_config = GenerationConfig(return_stats=True, format_version="v1")
+ results, _ = chatbot.predict_stream("Tell me about Intel Xeon Scalable Processors.", config=gen_config)
+ for text in results:
+ stream_text += text
+ print(text)
+ self.assertIn("END_OF_STREAM_STATS=", stream_text)
+
+ def test_text_chat_stream_return_stats(self):
+ config = PipelineConfig(model_name_or_path="facebook/opt-125m")
+ chatbot = build_chatbot(config)
+ stream_text = ""
+ gen_config = GenerationConfig(return_stats=True)
+ results, _ = chatbot.predict_stream("Tell me about Intel Xeon Scalable Processors.", config=gen_config)
+ for text in results:
+ stream_text += text
+ print(text)
+ self.assertIn("| Key | Value |", stream_text)
 
 if __name__ == '__main__':
  unittest.main()

workflows/chatbot/inference/generate.py

@@ -162,6 +162,12 @@ def parse_args():
  )
  parser.add_argument(
  "--return_stats", action='store_true', default=False,)
+ parser.add_argument(
+ "--format_version",
+ type=str,
+ default="v2",
+ help="the version of return stats format",
+ )
  args = parser.parse_args()
  return args
 
@@ -232,7 +238,9 @@ def main():
  use_hpu_graphs=args.use_hpu_graphs,
  use_cache=args.use_kv_cache,
  num_return_sequences=args.num_return_sequences,
- ipex_int8=args.ipex_int8
+ ipex_int8=args.ipex_int8,
+ return_stats=args.return_stats,
+ format_version=args.format_version
  )
 
  if args.habana: