spb/zyquo-mlx Public MIT

1//2//  DatasetService.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011enum DatasetImportError: LocalizedError {12    case emptyFile13    case undetectableFormat14    case allRowsInvalid(issues: [DatasetRowIssue])15    case tooFewSamples(count: Int)1617    var errorDescription: String? {18        switch self {19        case .emptyFile:20            "The file contains no data rows."21        case .undetectableFormat:22            "Could not detect the dataset format. Rows must be JSON objects with \"messages\", \"prompt\"+\"completion\", or \"text\" keys."23        case .allRowsInvalid(let issues):24            "No valid rows found. First problem (line \(issues.first?.line ?? 1)): \(issues.first?.problem ?? "")"25        case .tooFewSamples(let count):26            "Only \(count) valid samples — at least 8 are needed for a meaningful train/valid split."27        }28    }29}3031/// Result of a dataset import: what was written plus the validation report.32struct DatasetImportReport: Sendable {33    var dataset: Dataset34    var issues: [DatasetRowIssue]35    var skippedRows: Int36}3738/// Imports, validates, splits, and previews training datasets39/// (charter 3.A). mlx-lm does not auto-split local JSONL40/// (docs/TRAINING-RESEARCH.md §4.5) — the split is ours.41actor DatasetService {4243    static let shared = DatasetService()4445    private let root: URL4647    init(root: URL = PersistenceService.datasetsDirectory) {48        self.root = root49    }5051    // MARK: - Import5253    /// Import a JSONL file: validate every row, split into54    /// train/valid(/test), and write the dataset directory.55    func importJSONL(56        from sourceURL: URL,57        name: String,58        validFraction: Double = 0.1,59        testFraction: Double = 0.0,60        seed: UInt64 = 061    ) throws -> DatasetImportReport {62        let content = try String(contentsOf: sourceURL, encoding: .utf8)63        let lines = content.split(separator: "\n", omittingEmptySubsequences: true)64        guard !lines.isEmpty else { throw DatasetImportError.emptyFile }6566        var format: DatasetFormat?67        var validRows: [String] = []68        var issues: [DatasetRowIssue] = []69        var totalTokens = 070        var maxTokens = 07172        for (index, line) in lines.enumerated() {73            let lineNumber = index + 174            guard75                let data = line.data(using: .utf8),76                let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]77            else {78                issues.append(79                    DatasetRowIssue(80                        line: lineNumber, problem: "Not a valid JSON object.",81                        fix: "Each line must be one complete JSON object (no trailing commas, no multi-line objects)."))82                continue83            }8485            if format == nil {86                format = DatasetFormats.detect(firstRow: row)87                guard format != nil else { throw DatasetImportError.undetectableFormat }88            }8990            if let issue = DatasetFormats.validate(row: row, format: format!, line: lineNumber) {91                issues.append(issue)92                continue93            }9495            let sample = DatasetFormats.sample(from: row, format: format!)96            totalTokens += sample.estimatedTokens97            maxTokens = max(maxTokens, sample.estimatedTokens)98            validRows.append(String(line))99        }100101        guard let detectedFormat = format else { throw DatasetImportError.undetectableFormat }102        guard !validRows.isEmpty else { throw DatasetImportError.allRowsInvalid(issues: issues) }103        guard validRows.count >= 8 else { throw DatasetImportError.tooFewSamples(count: validRows.count) }104105        // Deterministic shuffle then split.106        var generator = SeededGenerator(seed: seed)107        validRows.shuffle(using: &generator)108109        let testCount = Int(Double(validRows.count) * testFraction)110        let validCount = max(1, Int(Double(validRows.count) * validFraction))111        let trainCount = validRows.count - validCount - testCount112113        let directory = root.appendingPathComponent(sanitize(name), isDirectory: true)114        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)115116        let train = validRows[0..<trainCount]117        let valid = validRows[trainCount..<(trainCount + validCount)]118        let test = validRows[(trainCount + validCount)...]119120        try write(rows: Array(train), to: directory.appendingPathComponent("train.jsonl"))121        try write(rows: Array(valid), to: directory.appendingPathComponent("valid.jsonl"))122        if !test.isEmpty {123            try write(rows: Array(test), to: directory.appendingPathComponent("test.jsonl"))124        }125126        let dataset = Dataset(127            id: UUID(),128            name: name,129            directory: directory,130            format: detectedFormat,131            trainCount: trainCount,132            validCount: validCount,133            testCount: test.count,134            totalTokens: totalTokens,135            maxSequenceTokens: maxTokens,136            importedAt: .now137        )138        try PersistenceService.saveJSON(dataset, to: directory.appendingPathComponent("dataset.json"))139140        return DatasetImportReport(141            dataset: dataset, issues: issues, skippedRows: issues.count)142    }143144    // MARK: - Library145146    func scan() throws -> [Dataset] {147        let fm = FileManager.default148        guard fm.fileExists(atPath: root.path) else { return [] }149        let entries = try fm.contentsOfDirectory(150            at: root, includingPropertiesForKeys: nil, options: .skipsHiddenFiles)151        return entries.compactMap {152            try? PersistenceService.loadJSON(153                Dataset.self, from: $0.appendingPathComponent("dataset.json"))154        }155        .sorted { $0.importedAt > $1.importedAt }156    }157158    /// First N samples of a split, parsed for preview.159    func preview(dataset: Dataset, split: String = "train", count: Int = 8) throws -> [DatasetSample] {160        let url = dataset.directory.appendingPathComponent("\(split).jsonl")161        let content = try String(contentsOf: url, encoding: .utf8)162        return content.split(separator: "\n").prefix(count).compactMap { line in163            guard164                let data = line.data(using: .utf8),165                let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]166            else { return nil }167            return DatasetFormats.sample(from: row, format: dataset.format)168        }169    }170171    func delete(_ dataset: Dataset) throws {172        try FileManager.default.removeItem(at: dataset.directory)173    }174175    // MARK: - Helpers176177    private func write(rows: [String], to url: URL) throws {178        try (rows.joined(separator: "\n") + "\n").write(to: url, atomically: true, encoding: .utf8)179    }180181    private func sanitize(_ name: String) -> String {182        name.replacingOccurrences(of: "[^A-Za-z0-9._-]+", with: "-", options: .regularExpression)183    }184}185186/// Deterministic RNG for reproducible splits (SplitMix64).187struct SeededGenerator: RandomNumberGenerator {188    private var state: UInt64189    init(seed: UInt64) { state = seed &+ 0x9E37_79B9_7F4A_7C15 }190    mutating func next() -> UInt64 {191        state &+= 0x9E37_79B9_7F4A_7C15192        var z = state193        z = (z ^ (z >> 30)) &* 0xBF58_476D_1CE4_E5B9194        z = (z ^ (z >> 27)) &* 0x94D0_49BB_1331_11EB195        return z ^ (z >> 31)196    }197}198