Repo layout

.github/workflows/check_code_quality.yml

@@ -0,0 +1,55 @@
+name: check_code_quality
+
+on:
+  push:
+    branches: [ main ]
+    paths:
+      - "setup.py"
+      - "optimum/tpu/**.py"
+      - "tests/**.py"
+      - "examples/**.py"
+
+  pull_request:
+    branches: [ main ]
+    paths:
+      - "setup.py"
+      - "optimum/tpu/**.py"
+      - "tests/**.py"
+      - "examples/**.py"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  build:
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: [3.10.12]
+        os: [ubuntu-22.04]
+
+    runs-on: ${{ matrix.os }}
+    steps:
+    - uses: actions/checkout@v4
+    - name: Setup Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Create and start a virtual environment
+      run: |
+        python -m venv venv
+        source venv/bin/activate
+    - name: Install dependencies
+      run: |
+        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
+        ruff check .

MANIFEST.in

@@ -0,0 +1,17 @@
+#  Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+include README.md
+include LICENSE
+

Makefile

@@ -1,8 +1,47 @@
+#  Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+SHELL := /bin/bash
+CURRENT_DIR = $(shell pwd)
+DEFAULT_CLONE_URL := https://github.com/huggingface/optimum-neuron.git
+# If CLONE_URL is empty, revert to DEFAULT_CLONE_URL
+REAL_CLONE_URL = $(if $(CLONE_URL),$(CLONE_URL),$(DEFAULT_CLONE_URL))
 
-VERSION := "0.0.1"
+.PHONY:	build_dist style style_check clean
 
 TGI_VERSION ?= 1.4.2
 
+rwildcard=$(wildcard $1) $(foreach d,$1,$(call rwildcard,$(addsuffix /$(notdir $d),$(wildcard $(dir $d)*))))
+
+VERSION := $(shell awk '/__version__ = "(.*)"/{print $$3}' optimum/tpu/version.py | sed 's/"//g')
+
+PACKAGE_DIST = dist/optimum-tpu-$(VERSION).tar.gz
+PACKAGE_WHEEL = dist/optimum_tpu-$(VERSION)-py3-none-any.whl
+PACKAGE_PYTHON_FILES = $(call rwildcard, optimum/*.py)
+PACKAGE_FILES = $(PACKAGE_PYTHON_FILES)  \
+				setup.py \
+				setup.cfg \
+				pyproject.toml \
+				README.md \
+				MANIFEST.in
+
+# Package build recipe
+$(PACKAGE_DIST) $(PACKAGE_WHEEL): $(PACKAGE_FILES)
+	python -m build
+
+clean:
+	rm -rf dist
+
 tpu-tgi:
 	docker build --rm -f text-generation-inference/Dockerfile \
 	             --build-arg VERSION=$(VERSION) \
@@ -24,19 +63,23 @@ build_dist_install_tools:
 	python -m pip install build
 	python -m pip install twine
 
+build_dist: ${PACKAGE_DIST} ${PACKAGE_WHEEL}
+
+pypi_upload: ${PACKAGE_DIST} ${PACKAGE_WHEEL}
+	python -m twine upload ${PACKAGE_DIST} ${PACKAGE_WHEEL}
 
+# Tests
 test_installs:
-	python -m pip install pytest safetensors
-	python -m pip install git+https://github.com/huggingface/transformers.git
+	python -m pip install .[tpu,tests]
+# 	python -m pip install git+https://github.com/huggingface/transformers.git
 
 # Stand-alone TGI server for unit tests outside of TGI container
 tgi_server:
 	python -m pip install -r text-generation-inference/server/build-requirements.txt
 	make -C text-generation-inference/server clean
 	VERSION=${VERSION} TGI_VERSION=${TGI_VERSION} make -C text-generation-inference/server gen-server
 
-tgi_test: tgi_server
-	python -m pip install pytest
+tgi_test: test_installs tgi_server
 	find text-generation-inference -name "text_generation_server-$(VERSION)-py3-none-any.whl" \
 	                               -exec python -m pip install --force-reinstall {} \;
 	python -m pytest -sv text-generation-inference/tests

optimum/tpu/__init__.py

@@ -0,0 +1,15 @@
+#  Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from .version import __version__, VERSION  # noqa: F401

optimum/tpu/generation/__init__.py

@@ -0,0 +1,17 @@
+# coding=utf-8
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .logits_process import FusedLogitsWarper  # noqa: F401
+from .token_selector import TokenSelector  # noqa: F401

text-generation-inference/server/text_generation_server/modeling.py → optimum/tpu/modeling.py

@@ -18,7 +18,6 @@
 from os import PathLike, environ
 from typing import Any
 
-import torch
 from loguru import logger
 from transformers import AutoModelForCausalLM
 from transformers.utils import is_accelerate_available
@@ -51,9 +50,7 @@ def 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 = AutoModelForCausalLM.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:

optimum/tpu/version.py

@@ -0,0 +1,18 @@
+#  Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+from pkg_resources import parse_version
+
+__version__ = "0.1.0.dev0"
+VERSION = parse_version(__version__)

pyproject.toml

@@ -11,6 +11,53 @@
 #  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
+[project]
+name = "optimum-tpu"
+dynamic = ["version"]
+authors=[
+    { name = "HuggingFace Inc. Machine Learning Optimization Team", email = "[email protected]"}
+]
+description = "Optimum TPU is the interface between the Hugging Face Transformers library and Google Cloud TPU devices."
+readme = "README.md"
+license = {file = "LICENSE"}
+classifiers = [
+    "Development Status :: 1 - Pre-Alpha",
+    "License :: OSI Approved :: Apache Software License",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Education",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+keywords = [
+    "transformers",
+    "fine-tuning",
+    "inference",
+    "tpu",
+    "cloud-tpu",
+    "gcp",
+    "google-cloud"
+]
+
+dependencies = ["transformers == 4.38.1", "torch-xla>=2.2.0", "torch>=2.2.0"]
+
+[project.optional-dependencies]
+tests = ["pytest", "safetensors"]
+quality = ["black", "ruff", "isort",]
+
+[project.urls]
+Homepage = "https://hf.co/hardware"
+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.packages.find]
+include = ["optimum.tpu"]
 
 [tool.black]
 line-length = 119
@@ -34,6 +81,4 @@ known-first-party = ["optimum.tpu"]
 [tool.pytest.ini_options]
 markers = [
     "is_staging_test",
-    "is_trainium_test",
-    "is_inferentia_test",
 ]

setup.cfg

@@ -0,0 +1,17 @@
+[isort]
+default_section = FIRSTPARTY
+ensure_newline_before_comments = True
+force_grid_wrap = 0
+include_trailing_comma = True
+known_first_party = optimum.tpu
+line_length = 119
+lines_after_imports = 2
+multi_line_output = 3
+use_parentheses = True
+
+[flake8]
+ignore = E203, E501, E741, W503, W605
+max-line-length = 119
+
+[tool:pytest]
+doctest_optionflags=NUMBER NORMALIZE_WHITESPACE ELLIPSIS