[NeuralChat] Fix response issue of model.predict (#1221)

* fix response issue of model.predict Signed-off-by: LetongHan <letong.han@intel.com>

Author: letonghan

intel_extension_for_transformers/neural_chat/models/model_utils.py

@@ -1445,18 +1445,21 @@ def predict(**params):
  output = tokenizer.decode(generation_output[0], skip_special_tokens=True)
  else:
  output = tokenizer.decode(generation_output.sequences[0], skip_special_tokens=True)
+
+ identifier_index = -1
  if "### Response:" in output:
- return output.split("### Response:")[-1].strip()
+ return output.split("### Response:")[identifier_index].strip()
  if "@@ Response" in output:
- return output.split("@@ Response")[-1].strip()
+ return output.split("@@ Response")[identifier_index].strip()
  if "### Assistant" in output:
- return output.split("### Assistant:")[-1].strip()
+ return output.split("### Assistant:")[identifier_index].strip()
  if "\nassistant\n" in output:
- return output.split("\nassistant\n")[-1].strip()
+ return output.split("\nassistant\n")[identifier_index].strip()
  if "[/INST]" in output:
- return output.split("[/INST]")[-1].strip()
+ return output.split("[/INST]")[identifier_index].strip()
  if "答：" in output:
- return output.split("答：")[-1].strip()
+ return output.split("答：")[identifier_index].strip()
  if "Answer:" in output:
- return output.split("Answer:")[-1].strip()
+ return output.split("Answer:")[identifier_index].strip()
+
  return output

intel_extension_for_transformers/neural_chat/tests/ci/models/test_model_utils.py

@@ -19,7 +19,7 @@
 import os
 import shutil
 from unittest import mock
-from intel_extension_for_transformers.neural_chat.models.model_utils import load_model, MODELS
+from intel_extension_for_transformers.neural_chat.models.model_utils import load_model, MODELS, predict
 from intel_extension_for_transformers.transformers import MixedPrecisionConfig, BitsAndBytesConfig, WeightOnlyQuantConfig
 from intel_extension_for_transformers.neural_chat.utils.common import get_device_type
 from intel_extension_for_transformers.neural_chat.utils.error_utils import clear_latest_error, get_latest_error
@@ -139,5 +139,19 @@ def test_model_optimization_weightonly(self):
  self.assertTrue("facebook/opt-125m" in MODELS)
  self.assertTrue(MODELS["facebook/opt-125m"]["model"] is not None)
 
+ @unittest.skipIf(get_device_type() != 'cpu', "Only run this test on CPU")
+ def test_model_predict(self):
+ load_model(model_name="facebook/opt-125m", tokenizer_name="facebook/opt-125m", device="cpu")
+ self.assertTrue("facebook/opt-125m" in MODELS)
+ self.assertTrue(MODELS["facebook/opt-125m"]["model"] is not None)
+
+ params = {
+ "model_name": "facebook/opt-125m",
+ "prompt": "hi"
+ }
+ output = predict(**params)
+ self.assertIn("hi", output)
+ self.assertNotIn("[/INST]", output)
+
 if __name__ == '__main__':
  unittest.main()