From ccbed88588d61fa9ed5a216a6242bd5c9e68ec4b Mon Sep 17 00:00:00 2001 From: Marina Zhang Date: Thu, 16 May 2024 17:43:15 +0000 Subject: [PATCH 1/5] change backend to default to onnx and install onnxruntime-gpu --- README.md | 4 +++- setup.py | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index d46ae36..cac0142 100644 --- a/README.md +++ b/README.md @@ -9,12 +9,14 @@ You can use `pip` to install the latest version of UniSim: pip install unisim ``` -By default, UniSim uses [Onnx](https://github.com/onnx/onnx) when running on CPU, and [TensorFlow](https://www.tensorflow.org/) for GPU acceleration. If you have a GPU, you can additionally install TensorFlow using: +By default, UniSim uses [Onnx](https://github.com/onnx/onnx) as the runtime. You can additionally install TensorFlow using: ``` pip install unisim[tensorflow] ``` +You can switch to using TensorFlow by setting the `BACKEND` environment variable (e.g. `os.environ["BACKEND"] = "tf"`). + ## Text UniSim (TextSim) The goal of TextSim is to provide an easy-to-use tool for efficient, accurate and multilingual fuzzy string matching, near-duplicate detection, and string similarity. Please see the tutorial [colab](notebooks/unisim_text_demo.ipynb) for an in-depth example on using TextSim for real-world use cases like fuzzy matching for addresses. diff --git a/setup.py b/setup.py index 41bb532..51ccf32 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,7 @@ def get_version(rel_path): author_email="unisim@google.com", url="https://github.com/google/unisim", license="MIT", - install_requires=["tabulate", "numpy", "tqdm", "onnx", "jaxtyping", "onnxruntime", "pandas", "usearch>=2.6.0"], + install_requires=["tabulate", "numpy", "tqdm", "onnx", "jaxtyping", "onnxruntime-gpu", "pandas", "usearch>=2.6.0"], extras_require={ "tensorflow": ["tensorflow>=2.11"], "dev": [ From 9c0fe4fee608b6674082d8765cd77ad5a17b7385 Mon Sep 17 00:00:00 2001 From: Marina Zhang Date: Thu, 16 May 2024 18:07:21 +0000 Subject: [PATCH 2/5] limit =2.6.0"], extras_require={ - "tensorflow": ["tensorflow>=2.11"], + "tensorflow": ["tensorflow>=2.11,<2.16"], "dev": [ "datasets", "mypy", From fd77b51dd4336f8860e306373479218827e6cfe0 Mon Sep 17 00:00:00 2001 From: Marina Zhang Date: Thu, 16 May 2024 23:31:49 +0000 Subject: [PATCH 3/5] add tf dependency to ensure the right version of TF is installed + backend cleanup --- .github/workflows/publish.yml | 2 +- .github/workflows/tests.yml | 2 +- README.md | 8 +------- setup.py | 3 +-- unisim/backend/load_backend.py | 18 +++++------------- 5 files changed, 9 insertions(+), 24 deletions(-) diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 5f0f06a..5d87f9b 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -25,7 +25,7 @@ jobs: - name: Install package run: | - pip install ".[tensorflow,dev]" + pip install ".[dev]" - name: Build package run: | diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 2c95607..5f5fefe 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -27,7 +27,7 @@ jobs: - name: Install package run: | - pip install ".[tensorflow,onnx,dev]" + pip install ".[dev]" - name: Lint with flake8 run: | diff --git a/README.md b/README.md index cac0142..8f87765 100644 --- a/README.md +++ b/README.md @@ -9,13 +9,7 @@ You can use `pip` to install the latest version of UniSim: pip install unisim ``` -By default, UniSim uses [Onnx](https://github.com/onnx/onnx) as the runtime. You can additionally install TensorFlow using: - -``` -pip install unisim[tensorflow] -``` - -You can switch to using TensorFlow by setting the `BACKEND` environment variable (e.g. `os.environ["BACKEND"] = "tf"`). +By default, UniSim uses [Onnx](https://github.com/onnx/onnx) as the runtime. You can switch to using TensorFlow by setting the `BACKEND` environment variable (e.g. `os.environ["BACKEND"] = "tf"`). ## Text UniSim (TextSim) diff --git a/setup.py b/setup.py index cf4aa3f..fe1958f 100644 --- a/setup.py +++ b/setup.py @@ -37,9 +37,8 @@ def get_version(rel_path): author_email="unisim@google.com", url="https://github.com/google/unisim", license="MIT", - install_requires=["tabulate", "numpy", "tqdm", "onnx", "jaxtyping", "onnxruntime-gpu", "pandas", "usearch>=2.6.0"], + install_requires=["tabulate", "numpy", "tqdm", "onnx", "jaxtyping", "onnxruntime-gpu", "pandas", "tensorflow>=2.11,<2.16", "usearch>=2.6.0"], extras_require={ - "tensorflow": ["tensorflow>=2.11,<2.16"], "dev": [ "datasets", "mypy", diff --git a/unisim/backend/load_backend.py b/unisim/backend/load_backend.py index 823f39f..9d9b4aa 100644 --- a/unisim/backend/load_backend.py +++ b/unisim/backend/load_backend.py @@ -30,18 +30,6 @@ except ImportError: TF_AVAILABLE = False -# detect accelerator -if TF_AVAILABLE or get_backend() == BackendType.tf: - devices_types = [d.device_type for d in tf.config.list_physical_devices()] - - if "GPU" in devices_types: - set_accelerator(AcceleratorType.gpu) - else: - set_accelerator(AcceleratorType.cpu) - -else: - set_accelerator(AcceleratorType.cpu) - # choose backend if not set by user accel = get_accelerator() backend = get_backend() @@ -63,9 +51,13 @@ # post detection if get_backend() == BackendType.onnx: from .onnx import * # noqa: F403, F401 + import onnxruntime as rt # FIXME onnx accelerator type support - set_accelerator(AcceleratorType.cpu) + if rt.get_device() == "GPU": + set_accelerator(AcceleratorType.gpu) + else: + set_accelerator(AcceleratorType.cpu) elif get_backend() == BackendType.tf: from .tf import * # type: ignore # noqa: F403, F401 From d7e14e9efdffac0f203280e8a7c58ea40bd148dd Mon Sep 17 00:00:00 2001 From: Marina Zhang Date: Thu, 16 May 2024 23:35:04 +0000 Subject: [PATCH 4/5] formatter --- setup.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index fe1958f..c31af56 100644 --- a/setup.py +++ b/setup.py @@ -37,7 +37,17 @@ def get_version(rel_path): author_email="unisim@google.com", url="https://github.com/google/unisim", license="MIT", - install_requires=["tabulate", "numpy", "tqdm", "onnx", "jaxtyping", "onnxruntime-gpu", "pandas", "tensorflow>=2.11,<2.16", "usearch>=2.6.0"], + install_requires=[ + "tabulate", + "numpy", + "tqdm", + "onnx", + "jaxtyping", + "onnxruntime-gpu", + "pandas", + "tensorflow>=2.11,<2.16", + "usearch>=2.6.0", + ], extras_require={ "dev": [ "datasets", From 5c48b656275ecbb0f39e901fe84f1afbfb29e4a1 Mon Sep 17 00:00:00 2001 From: Marina Zhang Date: Thu, 16 May 2024 23:39:48 +0000 Subject: [PATCH 5/5] fix imports --- unisim/backend/load_backend.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/unisim/backend/load_backend.py b/unisim/backend/load_backend.py index 9d9b4aa..6f6301b 100644 --- a/unisim/backend/load_backend.py +++ b/unisim/backend/load_backend.py @@ -50,9 +50,10 @@ # post detection if get_backend() == BackendType.onnx: - from .onnx import * # noqa: F403, F401 import onnxruntime as rt + from .onnx import * # noqa: F403, F401 + # FIXME onnx accelerator type support if rt.get_device() == "GPU": set_accelerator(AcceleratorType.gpu)