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'
+          python-version: '3.11'
-      - name: Install Poetry
+      - name: Install uv
         run: |
-          python -m pip install --upgrade pip
+          curl -LsSf https://astral.sh/uv/install.sh | sh
-          pip install poetry
+          echo "$HOME/.cargo/bin" >> $GITHUB_PATH
-      # 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
+      - name: Create virtual environment
-      # e.g. poetry version 0.1.${{ github.run_number }}
+        run: uv venv
-      - name: Set Version number
+      - name: Install build dependencies
         run: |
-          poetry version 1.3.0
+          uv pip install build ".[torch,flax]"
-      - name: Build and Publish to PyPI
+      - name: Build package
-        run: |
+        run: uv run python -m build
-          poetry config pypi-token.pypi ${{ secrets.PYPI_API_TOKEN }}
-          poetry publish --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
+2.  Create a virtual environment and install dependencies using `uv`:
-(CPU, GPU, TPU or Apple Silicon).
+    ```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>",
+    {name = "Rajat Sen", email = "senrajat@google.com"},
-    "Yichen Zhou <yichenzhou@google.com>",
+    {name = "Yichen Zhou", email = "yichenzhou@google.com"},
-    "Abhimanyu Das <abhidas@google.com>",
+    {name = "Abhimanyu Das", email = "abhidas@google.com"},
-    "Petros Mol <pmol@google.com>",
+    {name = "Petros Mol", email = "pmol@google.com"},
-    "Michael Chertushkin <chertushkinmichael@gmail.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]
+[project.optional-dependencies]
-python = ">=3.11,<4.0"
+torch = [
-numpy = ">=1.26.4"
+    "torch>=2.0.0",
-huggingface_hub = { version = ">=0.23.0", extras = ["cli"] }
+]
-safetensors = ">=0.5.3"
+flax = [
-torch = { version = ">=2.0.0", extras = ["cuda"], optional = true }
+    "flax",
-flax = { version = ">=0.8.2", optional = true }
+    "optax",
-optax = { version = ">=0.2.2", optional = true }
+    "einshape",
-einshape = { version = ">=0.8.0", optional = true }
+    "orbax-checkpoint",
-orbax-checkpoint = { version = ">=0.5.15", optional = true }
+    "jaxtyping",
-jaxtyping = { version = ">=0.2.29", optional = true }
+    "jax[cuda]"
-jax = { version = ">=0.4.26", optional = true }
+]
-
-
-[tool.poetry.extras]
-torch = ["torch"]
-flax = ["flax", "optax", "einshape", "orbax-checkpoint", "jaxtyping", "jax"]
 
 [tool.ruff]
 line-length = 88
 indent-width = 2
 
-
 [build-system]
-requires = ["poetry-core"]
+requires = ["setuptools>=61.0"]
-build-backend = "poetry.core.masonry.api"
+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