spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#pragma once34#include "core/tensor.h"56#include <cstdint>7#include <random>8#include <string>910namespace forge::train {1112// mmap'd uint16 token stream (llm.c format: 256 int32 header {20240520, 1,13// num_tokens}; headerless nanoGPT-style .bin also accepted). Batches are14// random contiguous windows of context_length+1; ids/targets are written15// into caller tensors as i32 (what the GPU kernels take).16class DataLoader {17public:18    DataLoader(const std::string& bin_path, int64_t context_length, uint64_t seed);19    ~DataLoader();2021    DataLoader(const DataLoader&) = delete;22    DataLoader& operator=(const DataLoader&) = delete;2324    // ids: [B, T] i32, targets: [B*T] i32 (targets = ids shifted by one)25    void next_batch(Tensor& ids, Tensor& targets);26    // Deterministic sequential window (evaluation); wraps around.27    void seq_batch(int64_t index, Tensor& ids, Tensor& targets) const;2829    int64_t num_tokens() const { return num_tokens_; }3031private:32    void fill(int64_t start, int64_t T, int32_t* ids_row, int32_t* tgt_row) const;3334    const uint16_t* tokens_ = nullptr; // into the mmap35    void* map_ = nullptr;36    size_t map_len_ = 0;37    int64_t num_tokens_ = 0;38    int64_t context_ = 0;39    std::mt19937_64 rng_;40};4142} // namespace forge::train43