From b0b7926930feaec9379599dbe1d8e8efa96e37e1 Mon Sep 17 00:00:00 2001 From: "Emiliano G.O." Date: Fri, 17 Jul 2026 14:22:28 -0300 Subject: [PATCH] Remove transitive datasets import from inference path Repoint two imports and package __init__.py so that 'import needle' and needle.model.run do not pull in the HuggingFace datasets library. All four symbols (get_tokenizer, to_snake_case, DEFAULT_MAX_ENC_LEN, DEFAULT_MAX_GEN_LEN) are defined in needle/dataset/tokenizer.py and were merely re-exported through needle/dataset/dataset.py. Changing the import source removes the unconditional from datasets import Dataset, load_from_disk at module level of dataset.py, which was the only thing tying the inference path to datasets (and its pyarrow/pandas/fsspec dependency tree). Needle/dataset/__init__.py also re-exported dataset.py via wildcard; removing that line is required because importing needle.dataset.tokenizer triggers package initialization, which previously loaded dataset.py. Signed-off-by: Emiliano Gandini Outeda --- needle/__init__.py | 2 +- needle/dataset/__init__.py | 1 - needle/model/run.py | 2 +- 3 files changed, 2 insertions(+), 3 deletions(-) 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,