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:
parent
1cc98b7007
commit
c2762f7f42
2 changed files with 56 additions and 1 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue