Add optional variants for tensor backends

jopetty · May 17, 2024 · 78d5077 · 78d5077
1 parent dcd4ca6
commit 78d5077
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -22,6 +22,18 @@ dependencies = [
 
 [project.optional-dependencies]
 dev = ["pre-commit"]
+torch = ["torch"]
+jax = [
+    "flax",
+    "jax"
+]
+hf = [
+    "accelerate",
+    "datasets",
+    "evaluate",
+    "transformers",
+    "tokenizers",
+]
 
 [tool.ruff]
 line-length = 88