markbaseengine/Sources/MarkBase/Tokenizer/Tokenizer.swift

import Foundation

// ─────────────────────────────────────────────────────────────
// Tokenizer Protocol - Unified interface for all tokenizers
// ─────────────────────────────────────────────────────────────

/// Tokenizer protocol for text-to-token and token-to-text conversion
public protocol Tokenizer: Sendable {
    /// Encode text to token IDs
    func encode(text: String) -> [Int]
    
    /// Decode token IDs to text
    func decode(tokens: [Int]) -> String
    
    /// Vocabulary size
    var vocabSize: Int { get }
    
    /// Special token IDs
    var bosTokenId: Int { get }
    var eosTokenId: Int { get }
    var eosTokenIds: Set<Int> { get }
    var padTokenId: Int { get }

    /// Raw token string for a given ID (used by StreamingDecoder)
    func rawToken(for id: Int) -> String?
}

// ─────────────────────────────────────────────────────────────
// Special Tokens Configuration
// ─────────────────────────────────────────────────────────────

public struct SpecialTokens: Sendable {
    public let bosToken: String
    public let eosToken: String
    public let padToken: String
    public let unkToken: String
    
    public init(
        bosToken: String = "<bos>",
        eosToken: String = "<eos>",
        padToken: String = "<pad>",
        unkToken: String = "<unk>"
    ) {
        self.bosToken = bosToken
        self.eosToken = eosToken
        self.padToken = padToken
        self.unkToken = unkToken
    }
    
    // Gemma-4 specific tokens
    public static let gemma4 = SpecialTokens(
        bosToken: "<bos>",
        eosToken: "<eos>",
        padToken: "<pad>",
        unkToken: "<unk>"
    )
}

// ─────────────────────────────────────────────────────────────
// Tokenizer Error
// ─────────────────────────────────────────────────────────────

public enum TokenizerError: Error, LocalizedError {
    case modelNotFound(String)
    case invalidModelFormat
    case encodingFailed(String)
    case decodingFailed([Int])
    case vocabMissing(String)
    
    public var errorDescription: String? {
        switch self {
        case .modelNotFound(let path):
            return "Tokenizer model not found at: \(path)"
        case .invalidModelFormat:
            return "Invalid tokenizer model format"
        case .encodingFailed(let text):
            return "Failed to encode text: \(text)"
        case .decodingFailed(let tokens):
            return "Failed to decode tokens: \(tokens)"
        case .vocabMissing(let token):
            return "Token not in vocabulary: \(token)"
        }
    }
}

// ─────────────────────────────────────────────────────────────
// Tokenizer Factory
// ─────────────────────────────────────────────────────────────

public final class TokenizerFactory: @unchecked Sendable {
    /// Load tokenizer from model directory
    public static func load(modelDir: String) throws -> Tokenizer {
        // Check tokenizer_config.json for tokenizer_class
        let configPath = modelDir + "/tokenizer_config.json"
        var tokenizerClass: String? = nil
        if FileManager.default.fileExists(atPath: configPath),
           let configData = try? Data(contentsOf: URL(fileURLWithPath: configPath)),
           let config = try? JSONSerialization.jsonObject(with: configData) as? [String: Any] {
            tokenizerClass = config["tokenizer_class"] as? String
        }
        
        // For GemmaTokenizer or SentencePiece-based tokenizers, prefer .model file
        if tokenizerClass == "GemmaTokenizer" || tokenizerClass?.contains("SentencePiece") == true {
            let modelPath = modelDir + "/tokenizer.model"
            if FileManager.default.fileExists(atPath: modelPath) {
                print("  Using SentencePieceTokenizer (tokenizer_class=\(tokenizerClass ?? "unknown"))")
                return try SentencePieceTokenizer(modelPath: modelPath)
            }
        }
        
        // Try tokenizer.json first (HuggingFace format)
        let tokenizerJsonPath = modelDir + "/tokenizer.json"
        if FileManager.default.fileExists(atPath: tokenizerJsonPath) {
            print("  Using BPETokenizer (tokenizer.json)")
            return try BPETokenizer(jsonPath: tokenizerJsonPath)
        }
        
        // Try .model file (SentencePiece format)
        let modelPath = modelDir + "/tokenizer.model"
        if FileManager.default.fileExists(atPath: modelPath) {
            print("  Using SentencePieceTokenizer (tokenizer.model)")
            return try SentencePieceTokenizer(modelPath: modelPath)
        }
        
        // Fallback to simple tokenizer (for testing)
        print("Warning: No tokenizer found, using simple tokenizer")
        return SimpleTokenizer()
    }
}