spb/forge-studio Public

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// A prepared dataset folder: train.bin/val.bin (llm.c-style header4// {magic 20240520, version, num_tokens} then uint16 tokens) + a forgebpe5// tokenizer model whose vocab size must match model.vocab_size.6import Foundation78struct Dataset: Identifiable, Equatable {9    var id: String { path }10    var path: String11    var name: String12    var trainTokens: Int?13    var valTokens: Int?14    var vocabSize: Int?15    var sizeOnDisk: Int641617    static let binMagic: Int32 = 20_240_5201819    // num_tokens lives at int32 offset 2 of the 256-int32 header.20    static func tokenCount(binURL: URL) -> Int? {21        guard let fh = try? FileHandle(forReadingFrom: binURL),22              let data = try? fh.read(upToCount: 12), data.count == 12 else {23            return nil24        }25        try? fh.close()26        let ints = data.withUnsafeBytes { $0.load(fromByteOffset: 0, as: (Int32, Int32, Int32).self) }27        guard ints.0 == binMagic else { return nil }28        return Int(ints.2)29    }3031    // forgebpe v1 header: line 1 "forgebpe v1", line 2 vocab size.32    static func vocabSize(tokenizerURL: URL) -> Int? {33        guard let fh = try? FileHandle(forReadingFrom: tokenizerURL),34              let data = try? fh.read(upToCount: 64),35              let head = String(data: data, encoding: .utf8) else { return nil }36        try? fh.close()37        let lines = head.split(separator: "\n", maxSplits: 2)38        guard lines.count >= 2, lines[0].hasPrefix("forgebpe") else { return nil }39        return Int(lines[1].trimmingCharacters(in: .whitespaces))40    }4142    static func scan(directory: URL) -> Dataset? {43        let fm = FileManager.default44        let train = directory.appendingPathComponent("train.bin")45        guard fm.fileExists(atPath: train.path) else { return nil }46        let val = directory.appendingPathComponent("val.bin")47        let tok = (try? fm.contentsOfDirectory(at: directory,48                                               includingPropertiesForKeys: nil))?49            .first { $0.pathExtension == "model" }50        var size: Int64 = 051        if let items = try? fm.contentsOfDirectory(at: directory,52                                                   includingPropertiesForKeys: [.fileSizeKey]) {53            for f in items {54                size += Int64((try? f.resourceValues(forKeys: [.fileSizeKey]).fileSize) ?? 0)55            }56        }57        return Dataset(58            path: directory.path,59            name: directory.lastPathComponent,60            trainTokens: tokenCount(binURL: train),61            valTokens: fm.fileExists(atPath: val.path) ? tokenCount(binURL: val) : nil,62            vocabSize: tok.flatMap { vocabSize(tokenizerURL: $0) },63            sizeOnDisk: size)64    }65}6667struct Checkpoint: Identifiable, Equatable {68    var id: String { path }69    var path: String70    var step: Int71    var sizeBytes: Int6472    var modifiedAt: Date73    var valLossAtStep: Double?74}75