From 352e19001c82b06832924055d66e7eb02e4ad582 Mon Sep 17 00:00:00 2001 From: Btara Truhandarien Date: Fri, 31 May 2024 15:14:19 -0400 Subject: [PATCH] fix(setup): install either tensorflow-macos or tensorflow depending on platform processor --- setup.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index ef80d06..11f28b9 100644 --- a/setup.py +++ b/setup.py @@ -26,6 +26,7 @@ import os from datetime import date from setuptools import find_packages, setup +import sys # We don't declare our dependency on transformers here because we build with # different packages for different variants @@ -61,7 +62,15 @@ extras["torch"] = ["torch>=1.8.0", "torchaudio"] # TODO: Remove upper bound of TF 2.11 once transformers release contains this fix: https://github.com/huggingface/evaluate/pull/372 -extras["tensorflow"] = ["tensorflow>=2.4.0,<2.11"] +if sys.platform == "darwin": + import platform + if platform.processor() == "arm": + tensorflow_versions = ["tensorflow-macos>=2.4.0,<2.11"] + else: + tensorflow_versions = ["tensorflow>=2.4.0,<2.11"] + extras["tensorflow"] = tensorflow_versions +else: + extras["tensorflow"] = ["tensorflow>=2.4.0,<2.11"] # MMS Server dependencies extras["mms"] = ["multi-model-server>=1.1.4", "retrying"]