File size: 1,715 Bytes
5e6d9f5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | """二进制 token 流的批采样。
语料存成一维 token 数组，每次随机截 seq_len+1 的窗口，
前 seq_len 个当输入、后 seq_len 个当标签。省内存、随机性够。
"""
import json
import os
from typing import Tuple
import numpy as np
import torch
class BinDataset:
def __init__(self, data_dir: str):
with open(os.path.join(data_dir, "meta.json"), "r", encoding="utf-8") as f:
self.meta = json.load(f)
self.dtype = np.dtype(self.meta["dtype"])
self.data_dir = data_dir
self._cache = {}
@property
def vocab_size(self) -> int:
return self.meta["vocab_size"]
def _arr(self, split: str) -> np.ndarray:
if split not in self._cache:
path = os.path.join(self.data_dir, f"{split}.bin")
self._cache[split] = np.memmap(path, dtype=self.dtype, mode="r")
return self._cache[split]
def get_batch(self, split: str, batch_size: int, seq_len: int,
device: torch.device, generator=None) -> Tuple[torch.Tensor, torch.Tensor]:
arr = self._arr(split)
hi = len(arr) - seq_len - 1
if hi <= 0:
raise ValueError(f"{split} 语料太短（{len(arr)} tokens），放不下 seq_len={seq_len}")
ix = torch.randint(hi, (batch_size,), generator=generator)
x = torch.stack([torch.from_numpy(arr[i:i + seq_len].astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy(arr[i + 1:i + 1 + seq_len].astype(np.int64)) for i in ix])
if device.type == "cuda":
return x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
return x.to(device), y.to(device)
|