spb/zyquo-mlx Public MIT

1//2//  DatasetFormats.swift3//  Zyquo MLX4//5//  Author: Simon-Pierre Boucher6//  Mail: contact@spboucher.ai7//89import Foundation1011/// A problem found in one dataset row, with a concrete fix12/// (charter 3.A: report malformed rows with fixes).13struct DatasetRowIssue: Identifiable, Codable, Hashable, Sendable {14    var id: Int { line }15    /// 1-based line number in the source file.16    var line: Int17    var problem: String18    var fix: String19}2021/// One parsed sample, normalized for preview.22struct DatasetSample: Sendable {23    var messages: [(role: String, content: String)]24    var estimatedTokens: Int25}2627/// Format detection + row validation for the JSONL shapes mlx-lm accepts28/// (docs/TRAINING-RESEARCH.md §4.2 — detection order: completions, chat, text).29enum DatasetFormats {3031    static func detect(firstRow row: [String: Any]) -> DatasetFormat? {32        if row["prompt"] != nil && row["completion"] != nil { return .completions }33        if row["messages"] != nil { return .chat }34        if row["text"] != nil { return .text }35        return nil36    }3738    /// Validate one row against a format. Returns an issue or nil when valid.39    static func validate(row: [String: Any], format: DatasetFormat, line: Int) -> DatasetRowIssue? {40        switch format {41        case .chat:42            guard let messages = row["messages"] as? [[String: Any]] else {43                return DatasetRowIssue(44                    line: line,45                    problem: "Missing or non-array \"messages\" key.",46                    fix: "Use {\"messages\": [{\"role\": \"user\", \"content\": \"…\"}, …]}.")47            }48            if messages.isEmpty {49                return DatasetRowIssue(50                    line: line, problem: "\"messages\" is empty.",51                    fix: "Provide at least a user and an assistant message.")52            }53            let validRoles: Set<String> = ["system", "user", "assistant", "tool"]54            for (i, message) in messages.enumerated() {55                guard let role = message["role"] as? String, validRoles.contains(role) else {56                    return DatasetRowIssue(57                        line: line,58                        problem: "Message \(i + 1) has a missing/invalid \"role\".",59                        fix: "Use one of: system, user, assistant, tool.")60                }61                guard message["content"] is String else {62                    return DatasetRowIssue(63                        line: line,64                        problem: "Message \(i + 1) (\(role)) has no string \"content\".",65                        fix: "Every message needs a \"content\" string.")66                }67            }68            if (messages.last?["role"] as? String) != "assistant" {69                return DatasetRowIssue(70                    line: line,71                    problem: "Conversation does not end with an assistant message.",72                    fix: "The final message must be the assistant reply the model should learn.")73            }74            return nil7576        case .completions:77            guard row["prompt"] is String else {78                return DatasetRowIssue(79                    line: line, problem: "Missing string \"prompt\".",80                    fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")81            }82            guard row["completion"] is String else {83                return DatasetRowIssue(84                    line: line, problem: "Missing string \"completion\".",85                    fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")86            }87            return nil8889        case .text:90            guard let text = row["text"] as? String, !text.isEmpty else {91                return DatasetRowIssue(92                    line: line, problem: "Missing or empty string \"text\".",93                    fix: "Use {\"text\": \"…\"} with non-empty content.")94            }95            return nil96        }97    }9899    /// Normalize a valid row into a preview sample.100    static func sample(from row: [String: Any], format: DatasetFormat) -> DatasetSample {101        var messages: [(String, String)] = []102        switch format {103        case .chat:104            for message in row["messages"] as? [[String: Any]] ?? [] {105                messages.append(106                    (message["role"] as? String ?? "?", message["content"] as? String ?? ""))107            }108        case .completions:109            messages = [110                ("user", row["prompt"] as? String ?? ""),111                ("assistant", row["completion"] as? String ?? ""),112            ]113        case .text:114            messages = [("text", row["text"] as? String ?? "")]115        }116        let characters = messages.reduce(0) { $0 + $1.1.count }117        // ~4 chars/token heuristic; precise counts come from the tokenizer at118        // training time (trainer truncates over max-seq-length with a warning).119        return DatasetSample(messages: messages, estimatedTokens: max(1, characters / 4))120    }121}122