spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "train/dataloader.h"34#include <fcntl.h>5#include <sys/mman.h>6#include <sys/stat.h>7#include <unistd.h>89#include <cstdio>10#include <cstdlib>11#include <cstring>1213namespace forge::train {1415namespace {16constexpr int32_t kMagic = 20240520;17constexpr size_t kHeaderInts = 256;1819[[noreturn]] void die(const std::string& msg) {20    std::fprintf(stderr, "forge/data: %s\n", msg.c_str());21    std::abort();22}23} // namespace2425DataLoader::DataLoader(const std::string& bin_path, int64_t context_length, uint64_t seed)26    : context_(context_length), rng_(seed) {27    const int fd = ::open(bin_path.c_str(), O_RDONLY);28    if (fd < 0) die("cannot open " + bin_path);29    struct stat st{};30    if (fstat(fd, &st) != 0) die("fstat failed on " + bin_path);31    map_len_ = size_t(st.st_size);32    map_ = mmap(nullptr, map_len_, PROT_READ, MAP_PRIVATE, fd, 0);33    ::close(fd);34    if (map_ == MAP_FAILED) die("mmap failed on " + bin_path);35    madvise(map_, map_len_, MADV_RANDOM);3637    const int32_t* header = static_cast<const int32_t*>(map_);38    if (map_len_ >= kHeaderInts * 4 && header[0] == kMagic && header[1] == 1) {39        num_tokens_ = header[2];40        tokens_ = reinterpret_cast<const uint16_t*>(header + kHeaderInts);41        if (size_t(num_tokens_) * 2 + kHeaderInts * 4 > map_len_)42            die("header token count exceeds file size: " + bin_path);43    } else {44        // headerless: the whole file is uint16 tokens45        num_tokens_ = int64_t(map_len_ / 2);46        tokens_ = static_cast<const uint16_t*>(map_);47    }48    if (num_tokens_ < context_ + 1)49        die("dataset smaller than one context window: " + bin_path);50}5152DataLoader::~DataLoader() {53    if (map_ && map_ != MAP_FAILED) munmap(map_, map_len_);54}5556void DataLoader::fill(int64_t start, int64_t T, int32_t* ids_row, int32_t* tgt_row) const {57    for (int64_t t = 0; t < T; ++t) {58        ids_row[t] = int32_t(tokens_[start + t]);59        tgt_row[t] = int32_t(tokens_[start + t + 1]);60    }61}6263void DataLoader::next_batch(Tensor& ids, Tensor& targets) {64    const int64_t B = ids.size(0), T = ids.size(1);65    std::uniform_int_distribution<int64_t> dist(0, num_tokens_ - T - 1);66    for (int64_t b = 0; b < B; ++b) {67        fill(dist(rng_), T, ids.data<int32_t>() + b * T, targets.data<int32_t>() + b * T);68    }69}7071void DataLoader::seq_batch(int64_t index, Tensor& ids, Tensor& targets) const {72    const int64_t B = ids.size(0), T = ids.size(1);73    const int64_t span = num_tokens_ - T - 1;74    for (int64_t b = 0; b < B; ++b) {75        const int64_t start = ((index * B + b) * T) % span;76        fill(start, T, ids.data<int32_t>() + b * T, targets.data<int32_t>() + b * T);77    }78}7980} // namespace forge::train81