diff --git a/needle/__init__.py b/needle/__init__.py
index f582872..f8a18d6 100644
--- a/needle/__init__.py
+++ b/needle/__init__.py
@@ -11,7 +11,7 @@
     encode_for_retrieval,
     retrieve_tools,
 )
-from needle.dataset.dataset import get_tokenizer
+from needle.dataset.tokenizer import get_tokenizer
 
 __all__ = [
     "SimpleAttentionNetwork",
diff --git a/needle/dataset/__init__.py b/needle/dataset/__init__.py
index 022e8de..94db9b7 100644
--- a/needle/dataset/__init__.py
+++ b/needle/dataset/__init__.py
@@ -1,2 +1 @@
 from .tokenizer import *  # noqa: F401,F403
-from .dataset import *  # noqa: F401,F403
diff --git a/needle/model/run.py b/needle/model/run.py
index 954022e..d9032f7 100644
--- a/needle/model/run.py
+++ b/needle/model/run.py
@@ -8,7 +8,7 @@
 import jax.numpy as jnp
 import numpy as np
 
-from ..dataset.dataset import get_tokenizer, to_snake_case, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_GEN_LEN
+from ..dataset.tokenizer import get_tokenizer, to_snake_case, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_GEN_LEN
 from .architecture import (
     SimpleAttentionNetwork,
     TransformerConfig,