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>
This commit is contained in:
Andrew Chen 2026-07-16 18:03:38 +08:00 committed by GitHub
commit c2762f7f42
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 56 additions and 1 deletions

View file

@ -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)

View file

@ -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)