diff --git a/angelslim/data/text_dataset.py b/angelslim/data/text_dataset.py index 54838a3a..5e9ae048 100644 --- a/angelslim/data/text_dataset.py +++ b/angelslim/data/text_dataset.py @@ -53,10 +53,6 @@ def _load_data(self, data_path: str, num_samples: int): def _load_hf_dataset(self, data_path: str, num_samples: int, block_size: int = 2048): parts = data_path.split(",") dataset = load_dataset(*parts)["train"] - total_samples = ( - min(num_samples, len(dataset["text"])) if num_samples > 0 else len(dataset["text"]) - ) - concatenated = {} for sample in dataset: tokenized = self.processor(sample["text"]) @@ -65,6 +61,9 @@ def _load_hf_dataset(self, data_path: str, num_samples: int, block_size: int = 2 concatenated[key] = [] concatenated[key].extend(tokenized[key]) + if not concatenated: + return + total = len(concatenated["input_ids"]) if total >= block_size: total = (total // block_size) * block_size @@ -72,13 +71,18 @@ def _load_hf_dataset(self, data_path: str, num_samples: int, block_size: int = 2 k: [t[i : i + block_size] for i in range(0, total, block_size)] for k, t in concatenated.items() } + total_samples = len(result["input_ids"]) + if num_samples > 0: + total_samples = min(num_samples, total_samples) for i in range(total_samples): inputs = { "input_ids": torch.tensor(result["input_ids"][i]).unsqueeze(0).to(self.device) } # HF CausalLM models shift labels internally; feed labels == input_ids. inputs["labels"] = inputs["input_ids"].clone() - inputs["attention_mask"] = torch.tensor(result["attention_mask"][i]).to(self.device) + inputs["attention_mask"] = ( + torch.tensor(result["attention_mask"][i]).unsqueeze(0).to(self.device) + ) self.data.append(inputs) def _load_parquet_data(self, data_path: str, num_samples: int): diff --git a/docs/source/design/prepare_dataset.md b/docs/source/design/prepare_dataset.md index bdb595b3..bb5b9044 100644 --- a/docs/source/design/prepare_dataset.md +++ b/docs/source/design/prepare_dataset.md @@ -47,6 +47,12 @@ jsonl格式为一行一条data，每行data为json格式。支持如下格式 [ { "from": "human", "value": "这本著作极其出色。请用同义词替换加粗的词语，构建一个新的句子。" }, { "from": "gpt", "value": "这部作品极为杰出。" } ] ``` +## Hugging Face文本数据集 + +`TextDataset`也支持将`data_path`设置为Hugging Face数据集标识，例如`Salesforce/wikitext,wikitext-2-raw-v1`，数据集需包含`train`分区和`text`列。文本会先拼接并按2048个token分块，`num_samples`限制的是分块后的校准样本数量；非正数表示使用所有可用块。 + +如果全部文本不足一个块，会保留这一个短样本；否则会丢弃最后不足一个块的尾部。空文本流不产生样本。 + ## DATA文件测试 ``` from transformers import AutoModelForCausalLM, AutoTokenizer @@ -67,4 +73,4 @@ dataloader = DataLoaderFactory.create_data_loader( num_samples=256, data_source=data_path, ) -``` \ No newline at end of file +``` diff --git a/tests/test_text_dataset_hf.py b/tests/test_text_dataset_hf.py new file mode 100644 index 00000000..3c7b76f4 --- /dev/null +++ b/tests/test_text_dataset_hf.py @@ -0,0 +1,80 @@ +# Copyright 2025 Tencent Inc. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Exercise HF text packing with CPU tensors and an offline tokenizer.""" + +import pytest +import torch +from datasets import Dataset, DatasetDict +from tokenizers import Tokenizer, models, pre_tokenizers +from torch.utils.data import DataLoader +from transformers import PreTrainedTokenizerFast + +from angelslim.data import text_dataset +from angelslim.data.base_dataset import BaseDataset + + +@pytest.fixture +def load_text_dataset(monkeypatch): + tokenizer = Tokenizer(models.WordLevel({"[UNK]": 0, "a": 1, "b": 2}, unk_token="[UNK]")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + processor = PreTrainedTokenizerFast(tokenizer_object=tokenizer, unk_token="[UNK]") + + def load(texts, num_samples=-1): + dataset = DatasetDict({"train": Dataset.from_dict({"text": texts})}) + monkeypatch.setattr(text_dataset, "load_dataset", lambda *args: dataset) + return text_dataset.TextDataset("offline-text-dataset", processor, num_samples=num_samples) + + return load + + +@pytest.mark.parametrize("num_samples", [-1, 0, 256]) +def test_short_rows_produce_one_partial_block(load_text_dataset, num_samples): + dataset = load_text_dataset(["a a", "b b"], num_samples) + + assert len(dataset) == 1 + assert dataset[0]["input_ids"].tolist() == [[1, 1, 2, 2]] + + +@pytest.mark.parametrize("num_samples,expected", [(-1, 3), (0, 3), (1, 1), (2, 2), (256, 3)]) +def test_sample_limit_counts_packed_blocks(load_text_dataset, num_samples, expected): + dataset = load_text_dataset(["a " * 6144], num_samples) + + assert len(dataset) == expected + assert all(item["input_ids"].shape == (1, 2048) for item in dataset) + + +def test_short_tail_after_full_blocks_is_still_dropped(load_text_dataset): + dataset = load_text_dataset(["a " * 2048, "b b"]) + + assert len(dataset) == 1 + assert dataset[0]["input_ids"].tolist() == [[1] * 2048] + + +@pytest.mark.parametrize("texts", [[], [""], ["", ""]]) +def test_empty_token_stream_produces_no_samples(load_text_dataset, texts): + assert len(load_text_dataset(texts)) == 0 + + +def test_attention_mask_batches_like_input_ids(load_text_dataset): + dataset = load_text_dataset(["a " * 2048, "b " * 2048], num_samples=2) + batch = next(iter(DataLoader(dataset, batch_size=2, collate_fn=BaseDataset.collate_fn))) + + assert batch["input_ids"].shape == (2, 2048) + assert batch["attention_mask"].shape == batch["input_ids"].shape + assert torch.equal(batch["attention_mask"], torch.ones_like(batch["input_ids"])) + assert torch.equal(batch["labels"], batch["input_ids"]) + assert batch["labels"].data_ptr() != batch["input_ids"].data_ptr() + assert batch["input_ids"][0].tolist() == [1] * 2048 + assert batch["input_ids"][1].tolist() == [2] * 2048