src/lmr/data/disk_dataset.py

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