Merge pull request #324 from google-research/rajat-dev

Use uv instead of poetry

.github/workflows/main.yml

@@ -1,4 +1,4 @@
-name: Poetry publish
+name: Python package build
 
 on:
   push:
@@ -7,24 +7,22 @@ on:
     branches: [ "master" ]
 
 jobs:
-  buildPush:
+  build:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: '3.10'
-      - name: Install Poetry
+          python-version: '3.11'
+      - name: Install uv
         run: |
-          python -m pip install --upgrade pip
-          pip install poetry
-      # use the hardcoded version number for now. if this is merged to the main repo, you can uncomment these lines to automatically tag the version number with the github run number
-      # e.g. poetry version 0.1.${{ github.run_number }}
-      - name: Set Version number
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+      - name: Create virtual environment
+        run: uv venv
+      - name: Install build dependencies
         run: |
-          poetry version 1.3.0
-      - name: Build and Publish to PyPI
-        run: |
-          poetry config pypi-token.pypi ${{ secrets.PYPI_API_TOKEN }}
-          poetry publish --build
+          uv pip install build ".[torch,flax]"
+      - name: Build package
+        run: uv run python -m build

.github/workflows/manual_publish.yml

@@ -0,0 +1,29 @@
+name: Manual PyPI Publish
+
+on:
+  workflow_dispatch:
+
+jobs:
+  build-and-publish:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: '3.11'
+      - name: Install uv
+        run: |
+          curl -LsSf https://astral.sh/uv/install.sh | sh
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
+      - name: Create virtual environment
+        run: uv venv
+      - name: Install build dependencies
+        run: uv pip install build twine
+      - name: Build package
+        run: uv run python -m build
+      - name: Publish to PyPI
+        env:
+          TWINE_USERNAME: __token__
+          TWINE_PASSWORD: ${{ secrets.PYPI_API_TOKEN }}
+        run: uv run twine upload dist/*

.gitignore

@@ -6,3 +6,4 @@ wandb/
 datasets/
 results/
 timesfm_jax.egg-info/
+development_setup.md

README.md

@@ -44,31 +44,33 @@ will be under construction over the next few weeks to
 
 ### Install
 
-**Step 1** (Prerequisite)
+1.  Clone the repository:
+    ```shell
+    git clone https://github.com/google-research/timesfm.git
+    cd timesfm
+    ```
 
-Install your preferred `torch` / `jax` backend based on your OS and accelerators
-(CPU, GPU, TPU or Apple Silicon).
+2.  Create a virtual environment and install dependencies using `uv`:
+    ```shell
+    # Create a virtual environment
+    uv venv
+    
+    # Activate the environment
+    source .venv/bin/activate
+    
+    # Install the package in editable mode with torch
+    uv pip install -e .[torch]
+    # Or with flax
+    uv pip install -e .[flax]
+    ```
+
+3. [Optional] Install your preferred `torch` / `jax` backend based on your OS and accelerators
+(CPU, GPU, TPU or Apple Silicon).:
 
 -   [Install PyTorch](https://pytorch.org/get-started/locally/).
 -   [Install Jax](https://docs.jax.dev/en/latest/installation.html#installation)
     for Flax.
 
-**Step 2** (Install `timesfm`)
-
-You can now install locally.
-
-```shell
-git clone https://github.com/google-research/timesfm.git
-cd timesfm
-pip install -e ".[torch]"
-```
-
-```shell
-git clone https://github.com/google-research/timesfm.git
-cd timesfm
-pip install -e ".[flax]"
-```
-
 ### Code Example
 
 ```python

pyproject.toml

@@ -1,42 +1,41 @@
-[tool.poetry]
+[project]
 name = "timesfm"
 version = "2.0.0"
 description = "A time series foundation model."
 authors = [
-    "Rajat Sen <senrajat@google.com>",
-    "Yichen Zhou <yichenzhou@google.com>",
-    "Abhimanyu Das <abhidas@google.com>",
-    "Petros Mol <pmol@google.com>",
-    "Michael Chertushkin <chertushkinmichael@gmail.com>"
+    {name = "Rajat Sen", email = "senrajat@google.com"},
+    {name = "Yichen Zhou", email = "yichenzhou@google.com"},
+    {name = "Abhimanyu Das", email = "abhidas@google.com"},
+    {name = "Petros Mol", email = "pmol@google.com"},
+    {name = "Michael Chertushkin", email = "chertushkinmichael@gmail.com"},
 ]
-license = "Apache-2.0"
+license = {text = "Apache-2.0"}
 readme = "README.md"
-packages = [{include = "timesfm", from = "src"}]
+requires-python = ">=3.11"
+dependencies = [
+    "numpy>=1.26.4",
+    "huggingface_hub[cli]>=0.23.0",
+    "safetensors>=0.5.3",
+]
 
-[tool.poetry.dependencies]
-python = ">=3.11,<4.0"
-numpy = ">=1.26.4"
-huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
-safetensors = ">=0.5.3"
-torch = { version = ">=2.0.0", extras = ["cuda"], optional = true }
-flax = { version = ">=0.8.2", optional = true }
-optax = { version = ">=0.2.2", optional = true }
-einshape = { version = ">=0.8.0", optional = true }
-orbax-checkpoint = { version = ">=0.5.15", optional = true }
-jaxtyping = { version = ">=0.2.29", optional = true }
-jax = { version = ">=0.4.26", optional = true }
-
-
-[tool.poetry.extras]
-torch = ["torch"]
-flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "jax"]
+[project.optional-dependencies]
+torch = [
+    "torch>=2.0.0",
+]
+flax = [
+    "flax",
+    "optax",
+    "einshape",
+    "orbax-checkpoint",
+    "jaxtyping",
+    "jax[cuda]"
+]
 
 [tool.ruff]
 line-length = 88
 indent-width = 2
 
-
 [build-system]
-requires = ["poetry-core"]
-build-backend = "poetry.core.masonry.api"
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
 

requirements.txt

@@ -0,0 +1,40 @@
+# This file was autogenerated by uv via the following command:
+#    uv pip compile pyproject.toml -o requirements.txt
+certifi==2025.10.5
+    # via requests
+charset-normalizer==3.4.3
+    # via requests
+filelock==3.19.1
+    # via huggingface-hub
+fsspec==2025.9.0
+    # via huggingface-hub
+hf-xet==1.1.10
+    # via huggingface-hub
+huggingface-hub==0.35.3
+    # via timesfm (pyproject.toml)
+idna==3.10
+    # via requests
+inquirerpy==0.3.4
+    # via huggingface-hub
+numpy==2.3.3
+    # via timesfm (pyproject.toml)
+packaging==25.0
+    # via huggingface-hub
+pfzy==0.3.4
+    # via inquirerpy
+prompt-toolkit==3.0.52
+    # via inquirerpy
+pyyaml==6.0.3
+    # via huggingface-hub
+requests==2.32.5
+    # via huggingface-hub
+safetensors==0.6.2
+    # via timesfm (pyproject.toml)
+tqdm==4.67.1
+    # via huggingface-hub
+typing-extensions==4.15.0
+    # via huggingface-hub
+urllib3==2.5.0
+    # via requests
+wcwidth==0.2.14
+    # via prompt-toolkit