From 9fd08237cbcd5f26476df60d2f297263c04c7ffe Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 10:04:57 +0200 Subject: [PATCH 1/9] Let's replace TPUModelForX with AutoModelForX to better match transformers --- optimum/tpu/modeling.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/optimum/tpu/modeling.py b/optimum/tpu/modeling.py index 5d6e7ac6..86b1945f 100644 --- a/optimum/tpu/modeling.py +++ b/optimum/tpu/modeling.py @@ -19,13 +19,13 @@ from typing import Any from loguru import logger -from transformers import AutoModelForCausalLM +from transformers import AutoModelForCausalLM as BaseAutoModelForCausalLM from transformers.utils import is_accelerate_available # TODO: For now TpuModelForCausalLM is just a shallow wrapper of # AutoModelForCausalLM, later this could be replaced by a custom class. -class TpuModelForCausalLM(AutoModelForCausalLM): +class AutoModelForCausalLM(BaseAutoModelForCausalLM): @classmethod def from_pretrained( @@ -46,11 +46,11 @@ def from_pretrained( else: device = "xla" if is_accelerate_available(): - model = AutoModelForCausalLM.from_pretrained( + model = BaseAutoModelForCausalLM.from_pretrained( pretrained_model_name_or_path, device_map=device, *model_args, **kwargs ) else: - model = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + model = BaseAutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) model.to(device) # Update config with specific data) if task is not None or getattr(model.config, "task", None) is None: From 2b673f2510f950743bdae22e6c79bc6d58a49e5b Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 10:07:45 +0200 Subject: [PATCH 2/9] Make AutoModelForX top-level import --- optimum/tpu/__init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/optimum/tpu/__init__.py b/optimum/tpu/__init__.py index adb37a59..bf203c2f 100644 --- a/optimum/tpu/__init__.py +++ b/optimum/tpu/__init__.py @@ -12,4 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. +from .modeling import AutoModelForCausalLM from .version import __version__, VERSION # noqa: F401 From 8df21f66fe9b577caaa637933abff8a95c7f8a62 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 10:07:57 +0200 Subject: [PATCH 3/9] match new names --- .../server/text_generation_server/generator.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/text-generation-inference/server/text_generation_server/generator.py b/text-generation-inference/server/text_generation_server/generator.py index 17fa439f..5df4bd4e 100644 --- a/text-generation-inference/server/text_generation_server/generator.py +++ b/text-generation-inference/server/text_generation_server/generator.py @@ -11,7 +11,7 @@ from loguru import logger from transformers import AutoTokenizer, PreTrainedTokenizerBase, StaticCache from transformers.generation import GenerationConfig -from optimum.tpu.modeling import TpuModelForCausalLM +from optimum.tpu import AutoModelForCausalLM from optimum.tpu.generation import TokenSelector from .pb.generate_pb2 import ( @@ -301,7 +301,7 @@ class TpuGenerator(Generator): def __init__( self, - model: TpuModelForCausalLM, + model, tokenizer: PreTrainedTokenizerBase, ): self.model = model @@ -633,7 +633,7 @@ def from_pretrained( """ logger.info("Loading model (this can take a few minutes).") start = time.time() - model = TpuModelForCausalLM.from_pretrained(model_path) + model = AutoModelForCausalLM.from_pretrained(model_path) end = time.time() logger.info(f"Model successfully loaded in {end - start:.2f} s.") tokenizer = AutoTokenizer.from_pretrained(model_path) From 63e2da57a92d91a5a1c1b05aad3e78f832976272 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 10:23:20 +0200 Subject: [PATCH 4/9] Remove tpu from make pip install --- Makefile | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Makefile b/Makefile index 5a0e0843..150ce469 100644 --- a/Makefile +++ b/Makefile @@ -70,7 +70,7 @@ pypi_upload: ${PACKAGE_DIST} ${PACKAGE_WHEEL} # Tests test_installs: - python -m pip install .[tpu,tests] + python -m pip install .[tests] tests: test_installs python -m pytest -sv tests From 43659524a5f5a72eaaead59cde4539562fc8086b Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 10:28:53 +0200 Subject: [PATCH 5/9] import version before anything else --- optimum/tpu/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/optimum/tpu/__init__.py b/optimum/tpu/__init__.py index bf203c2f..a0c43795 100644 --- a/optimum/tpu/__init__.py +++ b/optimum/tpu/__init__.py @@ -12,5 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .modeling import AutoModelForCausalLM from .version import __version__, VERSION # noqa: F401 +from .modeling import AutoModelForCausalLM From d5557ed455ddbc0dd18051736e6e09074b2608cd Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 11:15:39 +0200 Subject: [PATCH 6/9] Let's use setuptools_scm backend to install deps --- pyproject.toml | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index f5957a18..90675c21 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,10 @@ dependencies = [ "loguru == 0.6.0" ] +[build-system] +requires = ["setuptools>=64", "setuptools_scm>=8"] +build-backend = "setuptools.build_meta" + [project.optional-dependencies] tests = ["pytest", "safetensors"] quality = ["black", "ruff", "isort",] @@ -58,8 +62,8 @@ Documentation = "https://hf.co/docs/optimum/tpu" Repository = "https://github.com/huggingface/optimum-tpu" Issues = "https://github.com/huggingface/optimum-tpu/issues" -[tool.setuptools.dynamic] -version = {attr = "optimum.tpu.__version__"} +[tool.setuptools_scm] + [tool.setuptools.packages.find] include = ["optimum.tpu"] From b6af278c21e95c687e5e4b7f41783a0b0593c1e2 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 11:29:51 +0200 Subject: [PATCH 7/9] Fix F401 --- Makefile | 2 +- optimum/tpu/__init__.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Makefile b/Makefile index 150ce469..bbec4cb0 100644 --- a/Makefile +++ b/Makefile @@ -56,7 +56,7 @@ style_check: style: black . - ruff . --fix + ruff check . --fix # Utilities to release to PyPi build_dist_install_tools: diff --git a/optimum/tpu/__init__.py b/optimum/tpu/__init__.py index a0c43795..df32aeb3 100644 --- a/optimum/tpu/__init__.py +++ b/optimum/tpu/__init__.py @@ -12,5 +12,5 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .version import __version__, VERSION # noqa: F401 -from .modeling import AutoModelForCausalLM +from .version import __version__, VERSION # noqa: F401 +from .modeling import AutoModelForCausalLM # noqa: F401 From 7172845955de5db0c12619e20bd00956642344d1 Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 11:35:08 +0200 Subject: [PATCH 8/9] Let's just use ruff not black and ruff --- Makefile | 2 -- 1 file changed, 2 deletions(-) diff --git a/Makefile b/Makefile index bbec4cb0..82a4acea 100644 --- a/Makefile +++ b/Makefile @@ -51,11 +51,9 @@ tpu-tgi: # Run code quality checks style_check: - black --check . ruff . style: - black . ruff check . --fix # Utilities to release to PyPi From fd89de92cc90ef019890212162e5d98a20237f5f Mon Sep 17 00:00:00 2001 From: Morgan Funtowicz Date: Mon, 8 Apr 2024 11:39:43 +0200 Subject: [PATCH 9/9] Again --- .github/workflows/check_code_quality.yml | 4 ---- 1 file changed, 4 deletions(-) diff --git a/.github/workflows/check_code_quality.yml b/.github/workflows/check_code_quality.yml index fec93652..1ca2450b 100644 --- a/.github/workflows/check_code_quality.yml +++ b/.github/workflows/check_code_quality.yml @@ -45,10 +45,6 @@ jobs: source venv/bin/activate pip install --upgrade pip pip install .[quality] - - name: Check style with black - run: | - source venv/bin/activate - black --check . - name: Check style with ruff run: | source venv/bin/activate