Merge pull request #16 from google/v1.0.0

Version bump to 1.0.0  + small tweaks to backend

.github/workflows/tests.yml

@@ -27,7 +27,7 @@ jobs:
     
     - name: Install package
       run: |
-        pip install ".[dev]"
+        pip install ".[tensorflow,dev]"
         
     - name: Lint with flake8
       run: |
@@ -45,7 +45,7 @@ jobs:
     - name: Upload coverage to Coveralls
       if: ${{ github.ref == 'refs/heads/main' }}
       run: |
-        coveralls --service=github
+        coveralls -i --service=github
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
         COVERALLS_FLAG_NAME: ${{ matrix.python-version }}
@@ -74,6 +74,6 @@ jobs:
     - name: Finished
       if: ${{ github.ref == 'refs/heads/main' }}
       run: |
-        coveralls --service=github --finish
+        coveralls -i --service=github --finish
       env:
         GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

README.md

@@ -9,7 +9,11 @@ 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 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) when running on CPU, and [TensorFlow](https://www.tensorflow.org/) for GPU acceleration. You can switch backends by setting the `BACKEND` environment variable (e.g. `os.environ["BACKEND"] = "tf"` or `"onnx"`). If you have a GPU, you can additionally install TensorFlow using:
+2
+```
+pip install unisim[tensorflow]
+```
 
 ## Text UniSim (TextSim)
 

setup.py

@@ -41,14 +41,14 @@ def get_version(rel_path):
         "tabulate",
         "numpy",
         "tqdm",
-        "onnx",
         "jaxtyping",
-        "onnxruntime-gpu",
+        "onnx",
+        "onnxruntime",
         "pandas",
-        "tensorflow>=2.11,<2.16",
         "usearch>=2.6.0",
     ],
     extras_require={
+        "tensorflow": ["tensorflow>=2.11,<2.16"],
         "dev": [
             "datasets",
             "mypy",

unisim/__init__.py

@@ -4,5 +4,5 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
-__version__ = "0.0.2"
+__version__ = "1.0.0"
 from .textsim import TextSim  # noqa: F401

unisim/backend/load_backend.py

@@ -30,6 +30,18 @@
     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()
@@ -41,7 +53,7 @@
     elif accel == AcceleratorType.cpu:
         # on CPU always onnx
         set_backend(BackendType.onnx)
-    elif TF_AVAILABLE:
+    elif TF_AVAILABLE and accel == AcceleratorType.gpu:
         # on GPU use TF by default
         set_backend(BackendType.tf)
     else:
@@ -50,15 +62,10 @@
 
 # post detection
 if get_backend() == BackendType.onnx:
-    import onnxruntime as rt
-
     from .onnx import *  # noqa: F403, F401
 
-    # FIXME onnx accelerator type support
-    if rt.get_device() == "GPU":
-        set_accelerator(AcceleratorType.gpu)
-    else:
-        set_accelerator(AcceleratorType.cpu)
+    # FIXME(marinazh): onnx accelerator type support
+    set_accelerator(AcceleratorType.cpu)
 
 elif get_backend() == BackendType.tf:
     from .tf import *  # type: ignore # noqa: F403, F401

unisim/backend/tf.py

@@ -12,7 +12,6 @@
 from tensorflow import Tensor
 from tensorflow.keras import Model
 
-# typing
 from ..types import BatchEmbeddings
 
 

unisim/textsim.py

@@ -4,6 +4,8 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
+from __future__ import annotations
+
 from typing import Any, Dict, List, Sequence
 
 from pandas import DataFrame

unisim/types.py

@@ -4,15 +4,17 @@
 # license that can be found in the LICENSE file or at
 # https://opensource.org/licenses/MIT.
 
-from __future__ import annotations
-
-from typing import Union
+from typing import TYPE_CHECKING, Union
 
 from jaxtyping import Float32
-from numpy import ndarray
-from tensorflow import Tensor
 
-Array = Union[Tensor, ndarray]
+if TYPE_CHECKING:
+    from numpy import ndarray
+    from tensorflow import Tensor
+
+    Array = Union[Tensor, ndarray]
+else:
+    from numpy import ndarray as Array
 
 # Embeddings
 Embedding = Float32[Array, "embedding"]