from pathlib import Path import json import os from datetime import datetime from tqdm import tqdm import numpy as np import torch from torch.utils.data import Dataset class DiskDataset(Dataset): def __init__(self, file_path, max_seq_len, stride_fraction=None, allow_cycling=False): """ Args: file_path: path to the .bin file max_seq_len: - 在 stride_fraction=-1 (预填充模式) 下，必须与生成数据时的 pad_to_length 一致。 - 在普通模式下，表示窗口大小。 stride_fraction: - If > 0: Sliding window (standard pre-training). - If == -1: Pre-padded Sample Mode. Assumes data on disk is already padded to blocks of size `max_seq_len`. """ self.file_path = Path(file_path).resolve() assert self.file_path.is_file(), f"File not found: {self.file_path}" self.max_seq_len = max_seq_len self.stride_fraction = stride_fraction if stride_fraction is not None else 1.0 # import pdb # pdb.set_trace() # Determine mode self.pre_padded_mode = (self.stride_fraction == -1) self.data = np.memmap(self.file_path, dtype="int32", mode="r") self.file_size = len(self.data) if self.pre_padded_mode: # === Mode: Pre-padded Fixed Blocks === # 文件大小必须是 max_seq_len 的整数倍 if self.file_size % self.max_seq_len != 0: print(f"Warning: File size ({self.file_size}) is not a multiple of max_seq_len ({self.max_seq_len}). " f"Last partial sample might be ignored or dataset might be corrupted.") self.n_samples = self.file_size // self.max_seq_len self.stride = self.max_seq_len # Stride equals length in this mode else: # === Mode: Sliding Window === self.stride = int(self.max_seq_len * self.stride_fraction) self.n_samples = 1 + max(0, (self.file_size - self.max_seq_len) // self.stride) self.allow_cycling = allow_cycling and not self.pre_padded_mode def __len__(self): return self.n_samples def get_token_count(self): return self.file_size def __getitem__(self, idx): if self.allow_cycling: idx = idx % self.n_samples # 无论是预填充模式还是滑动窗口模式，读取逻辑其实是一样的： # 从 idx * stride 开始，读取 max_seq_len 长度 start = idx * self.stride end = start + self.max_seq_len # Create a copy to return a writeable tensor seq = np.array(self.data[start:end], dtype=np.int32, copy=True) return torch.from_numpy(seq).long() @staticmethod def generate_bin( dataset_iterator, tokenizer, output_path, add_eos=True, column="text", token_limit=None, metadata_path=None, pad_to_length=None, pad_token_id=-100, append: bool = False, # ✅ 新增 existing_metadata_path=None, # ✅ 可选：如果 append 时 metadata 分开存 ): output_path = Path(output_path) output_dir = output_path.parent output_dir.mkdir(parents=True, exist_ok=True) def tokenize_fn(text): ids = tokenizer.encode(text) if add_eos: ids.append(tokenizer.eos_token_id) return ids print(f"Building binary file at: {output_path}") if pad_to_length is not None: print(f"Mode: Fixed Block Size (Padded/Truncated to {pad_to_length})") else: print(f"Mode: Continuous Stream (Packed)") dtype = np.int32 bytes_per_token = np.dtype(dtype).itemsize # ====== ✅ append 支持：决定初始 pos、allocated ====== pos = 0 allocated = 0 mm = None # 如果 append=True 且文件存在，就从末尾继续写 if append and output_path.exists(): file_bytes = output_path.stat().st_size if file_bytes % bytes_per_token != 0: raise ValueError( f"Corrupted bin? File bytes {file_bytes} not divisible by token size {bytes_per_token}" ) pos = file_bytes // bytes_per_token allocated = pos # 先按当前大小映射，后面需要再 grow if allocated == 0: # 空文件当作新建 append = False else: mm = np.memmap(output_path, dtype=dtype, mode="r+", shape=(allocated,)) print(f"[INFO] Appending to existing file. Current tokens on disk: {pos}") else: # 原逻辑：覆盖重写 if output_path.exists(): output_path.unlink() # tqdm unit = "samples" pbar = tqdm(total=None, unit=unit) # Initial allocation size (in tokens) initial_alloc_tokens = 1_000_000 if pad_to_length: initial_alloc_tokens = (initial_alloc_tokens // pad_to_length) * pad_to_length def _grow(new_alloc_tokens): nonlocal mm, allocated if mm is not None: mm.flush() del mm mm = None new_bytes = new_alloc_tokens * bytes_per_token with open(output_path, "a+b") as f: if new_bytes > 0: f.seek(new_bytes - 1) f.write(b"\0") f.flush() os.fsync(f.fileno()) allocated = new_alloc_tokens mm = np.memmap(output_path, dtype=dtype, mode="r+", shape=(allocated,)) # ✅ 如果不是 append（或 append 但文件为空），创建初始空间 if mm is None: _grow(initial_alloc_tokens) else: # ✅ append 时：如果后续要写，仍需要预留增长空间 # 这里不马上 grow，等需要的时候再 grow pass # ====== ✅ 如果 append，要读入已有 metadata 以便累计 ====== prev_total_tokens = 0 prev_total_samples = 0 prev_pad_to_length = None meta_path_to_read = None if append: # 1) 优先从 existing_metadata_path 读取 if existing_metadata_path is not None and Path(existing_metadata_path).exists(): meta_path_to_read = Path(existing_metadata_path) # 2) 否则从 metadata_path 读取（如果你一直写同一个 metadata） elif metadata_path is not None and Path(metadata_path).exists(): meta_path_to_read = Path(metadata_path) if meta_path_to_read is not None: try: with open(meta_path_to_read, "r", encoding="utf-8") as f: prev = json.load(f) prev_total_tokens = int(prev.get("total_tokens", 0)) prev_total_samples = int(prev.get("total_samples", 0)) prev_pad_to_length = prev.get("pad_to_length", None) except Exception as e: print(f"[WARN] Failed to read previous metadata for append: {e}") # pad_to_length 一致性检查（很重要） if append and prev_pad_to_length != pad_to_length: raise ValueError( f"pad_to_length mismatch when appending: previous={prev_pad_to_length}, new={pad_to_length}" ) # ====== 主写入循环 ====== total_tokens_written = 0 sample_count = 0 done = False for example in dataset_iterator: if done: break ids = tokenize_fn(example[column]) if pad_to_length is not None: if len(ids) > pad_to_length: ids = ids[:pad_to_length] if len(ids) < pad_to_length: ids.extend([pad_token_id] * (pad_to_length - len(ids))) arr = np.asarray(ids, dtype=dtype) else: if not ids: continue arr = np.asarray(ids, dtype=dtype) # token_limit（如果你仍然想限制“本次追加写入”的 token 数） if token_limit is not None: if total_tokens_written + arr.size > token_limit: done = True if pad_to_length is None: remaining = token_limit - total_tokens_written arr = arr[:remaining] else: break needed = pos + arr.size if needed > allocated: new_alloc = max(max(allocated * 2, 1_000_000), needed) if pad_to_length: new_alloc = ((new_alloc + pad_to_length - 1) // pad_to_length) * pad_to_length _grow(new_alloc) mm[pos:pos + arr.size] = arr pos += arr.size total_tokens_written += arr.size sample_count += 1 pbar.update(1) pbar.close() if mm is not None: mm.flush() del mm mm = None # 截断到精确大小（append 时也一样需要） with open(output_path, "r+b") as f: f.truncate(pos * bytes_per_token) f.flush() os.fsync(f.fileno()) # ====== ✅ 写 metadata：累计总量 ====== if metadata_path is not None: meta = { "last_modified": datetime.now().isoformat(), "total_tokens": int(prev_total_tokens + total_tokens_written), "total_samples": int(prev_total_samples + sample_count), "pad_to_length": pad_to_length, "dtype": str(np.dtype(dtype)), "append": bool(append), } with open(metadata_path, "w", encoding="utf-8") as f: json.dump(meta, f, indent=2) print(f"✅ Wrote {total_tokens_written} tokens to {output_path} (int32)") if append: print(f"✅ Total tokens in file now ≈ {prev_total_tokens + total_tokens_written}") return total_tokens_written