diff --git a/python/llm/src/bigdl/llm/ggml/model/bloom/bloom_cpp.py b/python/llm/src/bigdl/llm/ggml/model/bloom/bloom_cpp.py index 18b5d65c..acde790b 100644 --- a/python/llm/src/bigdl/llm/ggml/model/bloom/bloom_cpp.py +++ b/python/llm/src/bigdl/llm/ggml/model/bloom/bloom_cpp.py @@ -64,31 +64,13 @@ from ctypes import ( c_size_t, ) import pathlib -from bigdl.llm.utils import get_avx_flags +from bigdl.llm.utils import get_shared_lib_info from bigdl.llm.utils.common import invalidInputError # Load the library def _load_shared_library(lib_base_name: str): - # Determine the file extension based on the platform - if sys.platform.startswith("linux") or sys.platform == "darwin": - lib_ext = ".so" - elif sys.platform == "win32": - lib_ext = ".dll" - else: - invalidInputError(False, "Unsupported platform") - - avx = get_avx_flags() - - # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs) - _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() - _base_path = _base_path / 'libs' - # Searching for the library in the current directory under the name "libbloom" (default name - # for bloomcpp) and "bloom" (default name for this repo) - _lib_paths = [ - _base_path / f"lib{lib_base_name}{avx}{lib_ext}", - _base_path / f"{lib_base_name}{avx}{lib_ext}", - ] + _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name) if "BLOOM_CPP_LIB" in os.environ: lib_base_name = os.environ["BLOOM_CPP_LIB"] diff --git a/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py b/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py index c1d2d2a9..30577624 100644 --- a/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py +++ b/python/llm/src/bigdl/llm/ggml/model/gptneox/gptneox_cpp.py @@ -64,30 +64,12 @@ from ctypes import ( ) import pathlib from bigdl.llm.utils.common import invalidInputError -from bigdl.llm.utils import get_avx_flags +from bigdl.llm.utils import get_shared_lib_info # Load the library def _load_shared_library(lib_base_name: str): - # Determine the file extension based on the platform - if sys.platform.startswith("linux") or sys.platform == "darwin": - lib_ext = ".so" - elif sys.platform == "win32": - lib_ext = ".dll" - else: - invalidInputError(False, "Unsupported platform.") - - avx = get_avx_flags() - - # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs) - _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() - _base_path = _base_path / 'libs' - # Searching for the library in the current directory under the name "libgptneox" (default name - # for gptneoxcpp) and "gptneox" (default name for this repo) - _lib_paths = [ - _base_path / f"lib{lib_base_name}{avx}{lib_ext}", - _base_path / f"{lib_base_name}{avx}{lib_ext}", - ] + _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name) if "GPTNEOX_CPP_LIB" in os.environ: lib_base_name = os.environ["GPTNEOX_CPP_LIB"] diff --git a/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py b/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py index 245f613a..68233c0f 100644 --- a/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py +++ b/python/llm/src/bigdl/llm/ggml/model/llama/llama_cpp.py @@ -64,30 +64,12 @@ from ctypes import ( ) import pathlib from bigdl.llm.utils.common import invalidInputError -from bigdl.llm.utils import get_avx_flags +from bigdl.llm.utils import get_shared_lib_info # Load the library def _load_shared_library(lib_base_name: str): - # Determine the file extension based on the platform - if sys.platform.startswith("linux") or sys.platform == "darwin": - lib_ext = ".so" - elif sys.platform == "win32": - lib_ext = ".dll" - else: - invalidInputError(False, "Unsupported platform.") - - avx = get_avx_flags() - - # Construct the paths to the possible shared library names - _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() - _base_path = _base_path / 'libs' - # Searching for the library in the current directory under the name "libllama" (default name - # for llamacpp) and "llama" (default name for this repo) - _lib_paths = [ - _base_path / f"lib{lib_base_name}{avx}{lib_ext}", - _base_path / f"{lib_base_name}{avx}{lib_ext}", - ] + _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name) if "LLAMA_CPP_LIB" in os.environ: lib_base_name = os.environ["LLAMA_CPP_LIB"] diff --git a/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder_cpp.py b/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder_cpp.py index 7de757db..f1b010b1 100644 --- a/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder_cpp.py +++ b/python/llm/src/bigdl/llm/ggml/model/starcoder/starcoder_cpp.py @@ -64,31 +64,13 @@ from ctypes import ( c_size_t, ) import pathlib -from bigdl.llm.utils import get_avx_flags +from bigdl.llm.utils import get_shared_lib_info from bigdl.llm.utils.common import invalidInputError # Load the library def _load_shared_library(lib_base_name: str): - # Determine the file extension based on the platform - if sys.platform.startswith("linux") or sys.platform == "darwin": - lib_ext = ".so" - elif sys.platform == "win32": - lib_ext = ".dll" - else: - invalidInputError(False, "Unsupported platform") - - avx = get_avx_flags() - - # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs) - _base_path = pathlib.Path(__file__).parent.parent.parent.parent.resolve() - _base_path = _base_path / 'libs' - # Searching for the library in the current directory under the name "libbloom" (default name - # for bloomcpp) and "bloom" (default name for this repo) - _lib_paths = [ - _base_path / f"lib{lib_base_name}{avx}{lib_ext}", - _base_path / f"{lib_base_name}{avx}{lib_ext}", - ] + _base_path, _lib_paths = get_shared_lib_info(lib_base_name=lib_base_name) if "STARCODER_CPP_LIB" in os.environ: lib_base_name = os.environ["STARCODER_CPP_LIB"] diff --git a/python/llm/src/bigdl/llm/utils/utils.py b/python/llm/src/bigdl/llm/utils/utils.py index f6064bca..996ff6d2 100644 --- a/python/llm/src/bigdl/llm/utils/utils.py +++ b/python/llm/src/bigdl/llm/utils/utils.py @@ -15,18 +15,44 @@ # import sys +import pathlib +from bigdl.llm.utils.isa_checker import check_avx_vnni, check_avx2, check_avx512_vnni from bigdl.llm.utils.common import invalidInputError, invalidOperationError -def get_avx_flags(): - avx = "" +def get_cpu_flags(): + flags = "" if sys.platform != "win32": - import subprocess - msg = subprocess.check_output(["lscpu"]).decode("utf-8") - if "avx512_vnni" in msg: - avx = "_avx512" - elif "avx2" in msg: - avx = "_avx2" + if check_avx512_vnni(): + flags = "_avx512" + elif check_avx_vnni(): + flags = "_avx2" else: invalidOperationError(False, "Unsupported CPUFLAGS.") - return avx + else: + flags = "_vnni" if check_avx_vnni() else "" + return flags + + +def get_shared_lib_info(lib_base_name: str): + # Determine the file extension based on the platform + if sys.platform.startswith("linux") or sys.platform == "darwin": + lib_ext = ".so" + elif sys.platform == "win32": + lib_ext = ".dll" + else: + invalidInputError(False, "Unsupported platform.") + + cpuflags = get_cpu_flags() + + # Construct the paths to the possible shared library names (python/llm/src/bigdl/llm/libs) + _base_path = pathlib.Path(__file__).parent.parent.resolve() + _base_path = _base_path / 'libs' + # Searching for the library in the current directory under the name "lib{lib_base_name}" + # (default name for llmcpp) and "{lib_base_name}" (default name for this repo) + _lib_paths = [ + _base_path / f"lib{lib_base_name}{cpuflags}{lib_ext}", + _base_path / f"{lib_base_name}{cpuflags}{lib_ext}", + ] + + return _base_path, _lib_paths diff --git a/python/llm/test/inference/test_call_models.py b/python/llm/test/inference/test_call_models.py index 04f53214..02bace3c 100644 --- a/python/llm/test/inference/test_call_models.py +++ b/python/llm/test/inference/test_call_models.py @@ -16,7 +16,6 @@ from bigdl.llm.models import Llama, Bloom, Gptneox, Starcoder -from bigdl.llm.utils import get_avx_flags import pytest from unittest import TestCase import os diff --git a/python/llm/test/langchain/test_langchain.py b/python/llm/test/langchain/test_langchain.py index b333fe64..e84923fc 100644 --- a/python/llm/test/langchain/test_langchain.py +++ b/python/llm/test/langchain/test_langchain.py @@ -14,7 +14,6 @@ # limitations under the License. # -from bigdl.llm.utils import get_avx_flags from bigdl.llm.langchain.embeddings import BigdlNativeEmbeddings from bigdl.llm.langchain.llms import BigdlNativeLLM import pytest