[Deps] Minimize deps (#617)

Author: zhiyuan1i

pyproject.toml

@@ -18,13 +18,12 @@ requires-python = ">=3.10"
 dependencies = [
     "torch>=2.5",
     "transformers>=4.53.0",
-    "datasets>=3.3.0",
     "einops",
 ]
 
 [project.optional-dependencies]
 conv1d = ["causal-conv1d>=1.4.0"]
-benchmark = ["matplotlib"]
+benchmark = ["matplotlib", "datasets>=3.3.0"]
 test = ["pytest"]
 
 [project.urls]

setup.py

@@ -41,12 +41,11 @@ def get_package_version():
     install_requires=[
         'torch>=2.5',
         'transformers>=4.53.0',
-        'datasets>=3.3.0',
         'einops'
     ],
     extras_require={
         'conv1d': ['causal-conv1d>=1.4.0'],
-        'benchmark': ['matplotlib'],
+        'benchmark': ['matplotlib', 'datasets>=3.3.0'],
         'test': ['pytest'],
     }
 )