fix(data): count Hugging Face samples after token packing

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,20 +61,28 @@ 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
         result = {
             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):

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,
 )
-```
+```

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