From c2762f7f421e53311bea8911767e632981d5cbdf Mon Sep 17 00:00:00 2001
From: Andrew Chen <48723787+chuenchen309@users.noreply.github.com>
Date: Thu, 16 Jul 2026 18:03:38 +0800
Subject: [PATCH] fix(dataprep): smart_chunk_text single-chunk path leaks
 internal tensor type when eos_token_id is None (#7151)

* fix(dataprep): smart_chunk_text single-chunk path leaks internal tensor type when eos_token_id is None

RawTextDataLoader.smart_chunk_text()'s single-chunk branch only
converts `tokens` to a plain Python list inside the
`if eos_token_id is not None:` guard. When a tokenizer has no
eos_token_id configured, that conversion is skipped entirely and the
function returns whatever internal tensor-like object came out of
the tokenizer normalization step (e.g. a torch.Tensor) as
"input_ids", instead of a list of ints.

The sibling multi-chunk branch a few lines below does the conversion
unconditionally, before checking eos_token_id -- the two branches of
the same method disagree on output type depending purely on whether
the tokenizer has an EOS token. Downstream, create_causal_dataset()
does `labels = [list(ids) for ids in input_ids]`; list()'ing a
tensor produces a list of 0-d tensor elements rather than plain
ints, inconsistent with every multi-chunk sample and liable to break
type inference in Dataset.from_dict()/downstream collation.

Fix: move the list conversion out of the eos_token_id guard,
matching the multi-chunk branch's existing pattern.

Added test_smart_chunk_text_single_chunk_no_eos_returns_plain_list
to tests/test_raw_text.py, confirmed red against unfixed code
(assertion failure: input_ids was a MockTensor, not a list) and
green after the fix. Full tests/test_raw_text.py (both test
functions) passes. ruff check + the repo's ruff-format-with-kwargs
script: clean.

Note: tests/test_raw_text.py does not appear to be wired into any
.github/workflows/*.yml CI job (a pre-existing repo characteristic,
not something introduced by this change) -- verified locally via
`python3 tests/test_raw_text.py`.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Etherl <61019402+Etherll@users.noreply.github.com>
---
 tests/test_raw_text.py       | 55 ++++++++++++++++++++++++++++++++++++
 unsloth/dataprep/raw_text.py |  2 +-
 2 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/tests/test_raw_text.py b/tests/test_raw_text.py
index c0f5d5f398..ba16e0cfc4 100644
--- a/tests/test_raw_text.py
+++ b/tests/test_raw_text.py
@@ -241,6 +241,61 @@ def test_raw_text_loader():
         os.unlink(test_file)
 
 
+def test_smart_chunk_text_single_chunk_no_eos_returns_plain_list():
+    """smart_chunk_text's single-chunk branch must return a plain list for
+    input_ids even when the tokenizer has no eos_token_id, matching the
+    multi-chunk branch's unconditional tolist()/list() conversion."""
+
+    class MockTensor:
+        def __init__(self, data):
+            self.data = data
+
+        def __getitem__(self, idx):
+            return self.data
+
+        def __len__(self):
+            return len(self.data)
+
+        def tolist(self):
+            return self.data
+
+    class MockTokenizerNoEos:
+        def __init__(self):
+            self.eos_token = None
+            self.eos_token_id = None
+
+        def __call__(
+            self,
+            text,
+            return_tensors = None,
+            add_special_tokens = False,
+        ):
+            token_ids = list(range(len(text.split())))
+            if return_tensors == "pt":
+                return {"input_ids": [MockTensor(token_ids)]}
+            return {"input_ids": token_ids}
+
+        def decode(
+            self,
+            token_ids,
+            skip_special_tokens = False,
+        ):
+            return " ".join(f"word_{i}" for i in token_ids)
+
+    loader = RawTextDataLoader(MockTokenizerNoEos(), chunk_size = 2048, stride = 512)
+    result = loader.smart_chunk_text(
+        "hello world short text", chunk_size = 2048, stride = 512, return_tokenized = True
+    )
+    input_ids = result[0]["input_ids"]
+    assert isinstance(
+        input_ids, list
+    ), f"input_ids should be a plain list even without an eos_token_id, got {type(input_ids)}"
+    assert input_ids == [0, 1, 2, 3], f"unexpected input_ids: {input_ids}"
+    print("✅ test_smart_chunk_text_single_chunk_no_eos_returns_plain_list passed!")
+    return True
+
+
 if __name__ == "__main__":
     success = test_raw_text_loader()
+    success = test_smart_chunk_text_single_chunk_no_eos_returns_plain_list() and success
     sys.exit(0 if success else 1)
diff --git a/unsloth/dataprep/raw_text.py b/unsloth/dataprep/raw_text.py
index 0993f0b4e9..128d966ecd 100644
--- a/unsloth/dataprep/raw_text.py
+++ b/unsloth/dataprep/raw_text.py
@@ -154,9 +154,9 @@ class RawTextDataLoader:
         if len(tokens) <= chunk_size:
             # Fits in a single chunk
             if return_tokenized:
+                tokens = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
                 eos_token_id = getattr(self.tokenizer, "eos_token_id", None)
                 if eos_token_id is not None:
-                    tokens = tokens.tolist() if hasattr(tokens, "tolist") else list(tokens)
                     tokens.append(eos_token_id)
 
                 attention_mask = [1] * len(tokens)