From 2266ca7d2b7cde4f52dec748c450ebbf739585ff Mon Sep 17 00:00:00 2001
From: Yuwen Hu <54161268+Oscilloscope98@users.noreply.github.com>
Date: Thu, 20 Jul 2023 13:20:25 +0800
Subject: [PATCH] [LLM] Small updates to transformers int4 ut (#8574)

* Small fix to transformers int4 ut

* Small fix
---
 .github/workflows/llm_unit_tests_linux.yml         |  1 -
 python/llm/test/inference/test_transformers_api.py |  6 ++----
 python/llm/test/run-llm-inference-tests.sh         | 10 ++++++----
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/.github/workflows/llm_unit_tests_linux.yml b/.github/workflows/llm_unit_tests_linux.yml
index 313931b5..266b9a12 100644
--- a/.github/workflows/llm_unit_tests_linux.yml
+++ b/.github/workflows/llm_unit_tests_linux.yml
@@ -76,7 +76,6 @@ jobs:
         uses: ./.github/actions/llm/setup-llm-env
         env:
           ANALYTICS_ZOO_ROOT: ${{ github.workspace }}
-          OMP_NUM_THREADS: 24
 
       - name: Download ckpt models
         run: |
diff --git a/python/llm/test/inference/test_transformers_api.py b/python/llm/test/inference/test_transformers_api.py
index f0f65fde..38130a9f 100644
--- a/python/llm/test/inference/test_transformers_api.py
+++ b/python/llm/test/inference/test_transformers_api.py
@@ -15,15 +15,13 @@
 #
 
 
-from bigdl.llm.models import Llama, Bloom, Gptneox, Starcoder
-from bigdl.llm.utils import get_avx_flags
 import unittest
 import os
 
 import time
 import torch
-from bigdl.llm.transformers import AutoModelForCausalLM, AutoModel
-from transformers import LlamaTokenizer, AutoTokenizer
+from bigdl.llm.transformers import AutoModel
+from transformers import AutoTokenizer
 
 class TestTransformersAPI(unittest.TestCase):
 
diff --git a/python/llm/test/run-llm-inference-tests.sh b/python/llm/test/run-llm-inference-tests.sh
index 81358339..816101fd 100644
--- a/python/llm/test/run-llm-inference-tests.sh
+++ b/python/llm/test/run-llm-inference-tests.sh
@@ -4,16 +4,18 @@ export ANALYTICS_ZOO_ROOT=${ANALYTICS_ZOO_ROOT}
 export LLM_HOME=${ANALYTICS_ZOO_ROOT}/python/llm/src
 export LLM_INFERENCE_TEST_DIR=${ANALYTICS_ZOO_ROOT}/python/llm/test/inference
 
-source bigdl-nano-init
-
 set -e
 
 echo "# Start testing inference"
 start=$(date "+%s")
 
 python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k "not test_transformers_int4"
-export OMP_NUM_THREADS=24
-taskset -c 0-23 python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k test_transformers_int4
+
+if [ -z "$THREAD_NUM" ]; then
+  THREAD_NUM=2
+fi
+export OMP_NUM_THREADS=$THREAD_NUM
+taskset -c 0-$((THREAD_NUM - 1)) python -m pytest -s ${LLM_INFERENCE_TEST_DIR} -k test_transformers_int4
 
 now=$(date "+%s")
 time=$((now-start))