[WIP] Add UT for Mistral Optimized Model (#9248)

* add ut for mistral model

* update

* fix model path

* upgrade transformers version for mistral model

* refactor correctness ut for mustral model

* refactor mistral correctness ut

* revert test_optimize_model back

* remove mistral from test_optimize_model

* add to revert transformers version back to 4.31.0

.github/workflows/llm_unit_tests.yml

@@ -83,6 +83,7 @@ jobs:
           echo "ORIGINAL_CHATGLM2_6B_PATH=${ORIGIN_DIR}/chatglm2-6b" >> "$GITHUB_ENV"
           echo "ORIGINAL_REPLIT_CODE_PATH=${ORIGIN_DIR}/replit-code-v1-3b" >> "$GITHUB_ENV"
           echo "ORIGINAL_WHISPER_TINY_PATH=${ORIGIN_DIR}/whisper-tiny" >> "$GITHUB_ENV"
+          echo "MISTRAL_ORIGIN_PATH=${ORIGIN_DIR}/Mistral-7B-v0.1" >> "$GITHUB_ENV"
 
           echo "LLAMA_INT4_CKPT_PATH=${INT4_CKPT_DIR}/bigdl_llm_llama_7b_q4_0.bin" >> "$GITHUB_ENV"
           echo "GPTNEOX_INT4_CKPT_PATH=${INT4_CKPT_DIR}/bigdl_llm_redpajama_7b_q4_0.bin" >> "$GITHUB_ENV"
@@ -146,6 +147,11 @@ jobs:
             echo "wget -r -nH --no-verbose --cut-dirs=1 $LLM_FTP_URL/llm/whisper-tiny -P $ORIGIN_DIR"
             wget -r -nH --no-verbose --cut-dirs=1 $LLM_FTP_URL/llm/whisper-tiny -P $ORIGIN_DIR
           fi
+          if [ ! -d $MISTRAL_ORIGIN_PATH ]; then
+            echo "Directory $MISTRAL_ORIGIN_PATH not found. Downloading from FTP server..."
+            echo "wget -r -nH --no-verbose --cut-dirs=1 $LLM_FTP_URL/llm/Mistral-7B-v0.1 -P $ORIGIN_DIR"
+            wget -r -nH --no-verbose --cut-dirs=1 $LLM_FTP_URL/llm/Mistral-7B-v0.1 -P $ORIGIN_DIR
+          fi
           if [ ! -d $LLAMA_ORIGIN_PATH ]; then
             echo "Directory $LLAMA_ORIGIN_PATH not found. Downloading from FTP server..."
             echo "wget --no-verbose $LLM_FTP_URL/llm/llama-7b-hf -P $ORIGIN_DIR"

python/llm/test/inference/test_optimize_mistral.py

@@ -0,0 +1,53 @@
+#
+# Copyright 2016 The BigDL Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import os
+import pytest
+
+from bigdl.llm.transformers import AutoModelForCausalLM
+from transformers import AutoTokenizer
+
+
+mistral_model_path = os.environ.get('MISTRAL_ORIGIN_PATH')
+
+prompt = "Once upon a time, there existed a little girl who liked to have adventures. She wanted to go to places and meet new people, and have fun"
+
+@pytest.mark.parametrize("Model, Tokenizer, model_path, prompt", [
+    (AutoModelForCausalLM, AutoTokenizer, mistral_model_path, prompt)
+])
+    
+def test_optimize_model(Model, Tokenizer, model_path, prompt):
+    tokenizer = Tokenizer.from_pretrained(model_path, trust_remote_code=True)
+    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+
+    model = Model.from_pretrained(model_path,
+                                load_in_4bit=True,
+                                optimize_model=False,
+                                trust_remote_code=True)
+    logits_base_model = (model(input_ids)).logits
+
+    model = Model.from_pretrained(model_path,
+                                load_in_4bit=True,
+                                optimize_model=True,
+                                trust_remote_code=True)
+    logits_optimized_model = (model(input_ids)).logits
+    diff = abs(logits_base_model - logits_optimized_model).flatten()
+
+    assert any(diff) is False
+
+
+if __name__ == '__main__':
+    pytest.main([__file__])

python/llm/test/inference/test_optimize_model.py

@@ -14,8 +14,8 @@
 # limitations under the License.
 #
 
-import pytest
 import os
+import pytest
 
 from bigdl.llm.transformers import AutoModelForCausalLM, AutoModel
 from transformers import LlamaTokenizer, AutoTokenizer
@@ -32,8 +32,9 @@ prompt = "Once upon a time, there existed a little girl who liked to have advent
     (AutoModelForCausalLM, LlamaTokenizer, llama_model_path, prompt),
     (AutoModelForCausalLM, AutoTokenizer, bloom_model_path, prompt),
     (AutoModel, AutoTokenizer, chatglm2_6b_model_path, prompt),
-    (AutoModelForCausalLM, AutoTokenizer, replit_code_model_path, prompt),
+    (AutoModelForCausalLM, AutoTokenizer, replit_code_model_path, prompt)
 ])
+    
 def test_optimize_model(Model, Tokenizer, model_path, prompt):
     tokenizer = Tokenizer.from_pretrained(model_path, trust_remote_code=True)
     input_ids = tokenizer.encode(prompt, return_tensors="pt")

python/llm/test/run-llm-inference-tests.sh

@@ -9,13 +9,19 @@ set -e
 echo "# Start testing inference"
 start=$(date "+%s")
 
-python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k "not test_transformers" -v
+python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k "not test_transformers" -v \
+      --ignore=${LLM_INFERENCE_TEST_DIR}/test_optimize_mistral.py
 
 if [ -z "$THREAD_NUM" ]; then
   THREAD_NUM=2
 fi
 export OMP_NUM_THREADS=$THREAD_NUM
-python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k test_transformers -v
+python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k test_transformers -v \
+      --ignore=${LLM_INFERENCE_TEST_DIR}/test_optimize_mistral.py
+
+python -m pip install transformers==4.34.0
+python -m pytest -s ${LLM_INFERENCE_TEST_DIR}/test_optimize_mistral.py -v
+python -m pip install transformers==4.31.0
 
 now=$(date "+%s")
 time=$((now-start))