Coverage for  / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / data / hf.py: 35%

1# Copyright 2026 Huawei Technologies Co., Ltd 

2# 

3# Licensed under the Apache License, Version 2.0 (the "License"); 

4# you may not use this file except in compliance with the License. 

5# You may obtain a copy of the License at 

6# 

7# http://www.apache.org/licenses/LICENSE-2.0 

8# 

9# Unless required by applicable law or agreed to in writing, software 

10# distributed under the License is distributed on an "AS IS" BASIS, 

11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 

12# See the License for the specific language governing permissions and 

13# limitations under the License. 

14# ============================================================================ 

15"""HuggingFace ``datasets`` builders (hub / local JSON). 

16 

17Two ``data.type`` aliases land here: 

18 

19- ``hf_datasets``: ``load_dataset(train_path, split="train")`` (hub or 

20 local arrow directory). Tokenises with the spec-provided transform. 

21- ``json_file``: ``load_dataset("json", data_files=train_path, split="train")`` 

22 — covers Alpaca-style ``instruction / input / output`` files. 

23 

24The wrapped :class:`TokenizedDataset` returns plain tensor rows so the 

25trainer's standard padding collator can stack them. ``max_steps`` 

26is auto-clipped to ``num_train_epochs * floor(len / global_batch_size)`` 

27so the dataloader never runs dry mid-epoch. 

28""" 

29import logging 

30from typing import Any 

31 

32import torch 

33from torch.utils.data import Dataset 

34 

35from hyper_parallel.data.registry import DATASET_REGISTRY 

36 

37 

38logger = logging.getLogger(__name__) 

39 

40 

41class TokenizedDataset(Dataset): 

42 """Lightweight ``torch.utils.data.Dataset`` view over an HF dataset. 

43 

44 Each row is converted to ``torch.long`` tensors on demand — the HF 

45 table stays on disk / memory in arrow format until the dataloader 

46 actually pulls a sample. 

47 """ 

48 

49 def __init__(self, hf_ds: Any) -> None: 

50 self.data = hf_ds 

51 

52 def __len__(self) -> int: 

53 return len(self.data) 

54 

55 def __getitem__(self, idx: int): 

56 item = self.data[idx] 

57 return { 

58 "input_ids": torch.tensor(item["input_ids"], dtype=torch.long), 

59 "labels": torch.tensor(item["labels"], dtype=torch.long), 

60 } 

61 

62 

63def _load_raw(args: Any, data_type: str) -> Any: 

64 """Run the appropriate ``load_dataset`` call for ``data_type``.""" 

65 from datasets import load_dataset # pylint: disable=C0415 # optional dep 

66 

67 train_path = args.data.train_path 

68 if not train_path: 

69 raise ValueError(f"data.train_path is required when data.type='{data_type}'") 

70 

71 if data_type == "json_file": 

72 return load_dataset("json", data_files=train_path, split="train") 

73 subset = args.data.subset 

74 if subset: 

75 return load_dataset(train_path, subset, split="train") 

76 return load_dataset(train_path, split="train") 

77 

78 

79def _maybe_truncate(ds: Any, args: Any) -> Any: 

80 """Apply ``data.train_size`` if it shrinks the dataset.""" 

81 train_size = args.data.train_size 

82 if train_size and train_size < len(ds): 

83 ds = ds.select(range(train_size)) 

84 logger.info("Dataset truncated to %d samples", train_size) 

85 return ds 

86 

87 

88def _build_hf(*, base: Any, args: Any, data_transform: Any, data_type: str, **_: Any) -> TokenizedDataset: 

89 """Shared loader for ``hf_datasets`` and ``json_file``. 

90 

91 Tokenises via ``data_transform`` when provided, then filters away empty 

92 sequences and wraps in :class:`TokenizedDataset`. Updates 

93 ``base.state.max_steps`` to ``min(cfg.max_steps, len/global_bs)`` so 

94 epoch boundaries stay consistent with the data on hand. 

95 """ 

96 logger.info( 

97 "Loading dataset: type=%s, path=%s", data_type, args.data.train_path, 

98 ) 

99 ds = _maybe_truncate(_load_raw(args, data_type), args) 

100 

101 if data_transform is not None: 

102 ds = ds.map( 

103 data_transform, 

104 batched=True, 

105 remove_columns=ds.column_names, 

106 desc="Tokenizing", 

107 ) 

108 # Drop empty rows only when the dataset is already tokenized — without a 

109 # transform a raw text dataset has no ``input_ids`` column and the filter 

110 # predicate would raise ``KeyError`` instead of loading. 

111 if "input_ids" in ds.column_names: 

112 ds = ds.filter(lambda x: len(x["input_ids"]) > 0) 

113 

114 wrapped = TokenizedDataset(ds) 

115 # max_steps clipping must happen pre-dataloader so the train loop's 

116 # epoch count matches the data on hand; the trainer reads 

117 # ``base.state.max_steps`` further down the build chain. The cap is 

118 # the TOTAL step budget across all epochs — clipping to a single 

119 # epoch would silently truncate multi-epoch training. 

120 num_epochs = max(int(args.train.num_train_epochs or 1), 1) 

121 base.state.max_steps = min( 

122 args.train.max_steps, 

123 num_epochs * (len(wrapped) // max(args.train.global_batch_size, 1)), 

124 ) 

125 logger.info( 

126 "Dataset ready: %d samples, max_steps=%d", 

127 len(wrapped), base.state.max_steps, 

128 ) 

129 return wrapped 

130 

131 

132@DATASET_REGISTRY.register("hf_datasets") 

133def build_hf_datasets(**kwargs: Any) -> TokenizedDataset: 

134 """Build an HF hub / arrow dataset.""" 

135 return _build_hf(data_type="hf_datasets", **kwargs) 

136 

137 

138@DATASET_REGISTRY.register("json_file") 

139def build_json_file(**kwargs: Any) -> TokenizedDataset: 

140 """Build a local Alpaca-style ``.json`` / ``.jsonl`` dataset.""" 

141 return _build_hf(data_type="json_file", **kwargs)