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)