Merge pull request #16 from tubo213:feature/optional-dependencies

Feature/optional-dependencies

README.md

@@ -21,6 +21,18 @@ $ pip install uv
 uv sync
 ```
 
+if you need to install Pytorch(cpu) and torchvision, run the following command.
+
+```
+uv sync --extra cpu
+```
+
+if you need to install Pytorch(cu121) and torchvision, run the following command.
+
+```
+uv sync --extra cu121
+```
+
 3. activate pre-commit
 
 ```

pyproject.toml

@@ -9,31 +9,24 @@ dependencies = [
     "hydra-core>=1.3.2",
     "scipy>=1.10.1",
     "pandas>=2.0.3",
-    "polars[gpu]>=1.5.0",
+    "polars>=1.5.0",
     "scikit-learn>=1.3.0",
     "loguru>=0.7.2",
     "tqdm>=4.64.1",
-    "lightning==2.4.0",
     "matplotlib>=3.8.4",
     "seaborn>=0.13.2",
     "japanize_matplotlib>=1.1.3",
     "wandb>=0.16.6",
-    "torch>=2.4.1",
 ]
 readme = "README.md"
 requires-python = ">= 3.10"
 
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
-[tool.hatch.build.targets.wheel]
-packages = ["src"]
-
-[tool.uv.pip]
-index-url = "https://download.pytorch.org/whl/cu121"
+[project.optional-dependencies]
+cpu = ["torch>=2.5.1", "torchvision>=0.20.1", "lightning==2.4.0"]
+cu121 = ["torch>=2.5.1", "torchvision>=0.20.1", "lightning==2.4.0"]
 
 [tool.uv]
+conflicts = [[{ extra = "cpu" }, { extra = "cu121" }]]
 managed = true
 dev-dependencies = [
     "jupyter>=1.1.1",
@@ -42,6 +35,33 @@ dev-dependencies = [
     "ruff>=0.6.9",
 ]
 
+[tool.uv.sources]
+torch = [
+    { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
+    { index = "pytorch-cu121", extra = "cu121" },
+]
+torchvision = [
+    { index = "pytorch-cpu", extra = "cpu", marker = "platform_system != 'Darwin'" },
+    { index = "pytorch-cu121", extra = "cu121" },
+]
+
+[[tool.uv.index]]
+name = "pytorch-cpu"
+url = "https://download.pytorch.org/whl/cpu"
+explicit = true
+
+[[tool.uv.index]]
+name = "pytorch-cu121"
+url = "https://download.pytorch.org/whl/cu121"
+explicit = true
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src"]
+
 [tool.mypy]
 allow_redefinition = true
 allow_untyped_globals = false