mlboydaisuke/Granite-Embedding-97M-Multilingual-R2-CoreAI

1import Foundation2 3/// Independent tokenizer for the pinned Granite 97M R2 raw-text contract.4/// No prompts, trimming, Unicode normalization, HF package, or model execution.5final class GraniteTokenizer {6    enum TokenizerError: Error, CustomStringConvertible {7        case unsupported(String)8        case missingToken(String)9        var description: String {10            switch self {11            case .unsupported(let message): return "Unsupported tokenizer: \(message)"12            case .missingToken(let token): return "Missing BPE token: \(token)"13            }14        }15    }16 17    private struct AddedToken: Decodable {18        let id: Int3219        let content: String20        let single_word: Bool21        let lstrip: Bool22        let rstrip: Bool23        let normalized: Bool24    }25    private struct Model: Decodable {26        let type: String27        let ignore_merges: Bool28        let byte_fallback: Bool29        let dropout: Double?30        let unk_token: String?31        let continuing_subword_prefix: String?32        let end_of_word_suffix: String?33        let vocab: [String: Int32]34        let merges: [[String]]35    }36    private struct TokenizerData: Decodable {37        let model: Model38        let added_tokens: [AddedToken]39    }40    private let vocab: [String: Int32]41    private let ranks: [String: Int]42    private let pattern: NSRegularExpression43    private let addedPattern: NSRegularExpression44    private let wordStart: NSRegularExpression45    private let wordEnd: NSRegularExpression46    private let added: [String: AddedToken]47    private let byteAlphabet: [String]48    private var cache: [String: [Int32]] = [:]49    private let clsID: Int32 = 17993450    private let sepID: Int32 = 17993851    private let padID: Int32 = 17993552 53    init(tokenizerURL: URL, tokenizerConfigURL: URL) throws {54        let bytes = try Data(contentsOf: tokenizerURL)55        let typed = try JSONDecoder().decode(TokenizerData.self, from: bytes)56        guard let root = try JSONSerialization.jsonObject(with: bytes) as? [String: Any],57              root["normalizer"] is NSNull,58              let pre = root["pre_tokenizer"] as? [String: Any],59              pre["type"] as? String == "Sequence",60              let steps = pre["pretokenizers"] as? [[String: Any]], steps.count == 2,61              steps[0]["type"] as? String == "Split",62              steps[0]["behavior"] as? String == "Isolated",63              steps[0]["invert"] as? Bool == false,64              let regexPattern = steps[0]["pattern"] as? [String: String],65              regexPattern.count == 1, let expression = regexPattern["Regex"],66              steps[1]["type"] as? String == "ByteLevel",67              steps[1]["add_prefix_space"] as? Bool == false,68              steps[1]["use_regex"] as? Bool == false69        else { throw TokenizerError.unsupported("normalizer or pretokenizer") }70        let model = typed.model71        guard model.type == "BPE", model.ignore_merges, !model.byte_fallback,72              model.dropout == nil, model.unk_token == nil,73              model.continuing_subword_prefix == nil, model.end_of_word_suffix == nil74        else { throw TokenizerError.unsupported("BPE configuration") }75        guard let post = root["post_processor"] as? [String: Any],76              post["type"] as? String == "TemplateProcessing",77              let single = post["single"] as? [[String: Any]], single.count == 3,78              (single[0]["SpecialToken"] as? [String: Any])?["id"] as? String == "<|startoftext|>",79              (single[1]["Sequence"] as? [String: Any])?["id"] as? String == "A",80              (single[2]["SpecialToken"] as? [String: Any])?["id"] as? String == "<|return|>"81        else { throw TokenizerError.unsupported("single-text template") }82        let configBytes = try Data(contentsOf: tokenizerConfigURL)83        guard let config = try JSONSerialization.jsonObject(with: configBytes) as? [String: Any],84              (config["padding_side"] as? String ?? "right") == "right",85              (config["truncation_side"] as? String ?? "right") == "right",86              config["cls_token"] as? String == "<|startoftext|>",87              config["sep_token"] as? String == "<|return|>",88              config["pad_token"] as? String == "<|endoftext|>"89        else { throw TokenizerError.unsupported("special tokens/padding/truncation") }90 91        var added: [String: AddedToken] = [:]92        for token in typed.added_tokens {93            guard !token.lstrip, !token.rstrip, !token.normalized, !token.content.isEmpty94            else { throw TokenizerError.unsupported("added-token flags") }95            added[token.content] = token96        }97        guard added["<|startoftext|>"]?.id == 179934,98              added["<|return|>"]?.id == 179938,99              added["<|endoftext|>"]?.id == 179935100        else { throw TokenizerError.unsupported("special IDs") }101        self.added = added102        self.vocab = model.vocab103        var ranks: [String: Int] = [:]104        ranks.reserveCapacity(model.merges.count)105        for (rank, pair) in model.merges.enumerated() {106            guard pair.count == 2, !pair[0].contains("\0"), !pair[1].contains("\0")107            else { throw TokenizerError.unsupported("merge pair") }108            ranks[pair[0] + "\0" + pair[1]] = rank109        }110        self.ranks = ranks111        self.pattern = try NSRegularExpression(pattern: expression)112        let addedExpression = added.keys.sorted {113            $0.utf8.count == $1.utf8.count ? $0 < $1 : $0.utf8.count > $1.utf8.count114        }.map(NSRegularExpression.escapedPattern(for:)).joined(separator: "|")115        self.addedPattern = try NSRegularExpression(pattern: addedExpression)116        self.wordStart = try NSRegularExpression(pattern: #"^\w"#)117        self.wordEnd = try NSRegularExpression(pattern: #"\w$"#)118 119        let visible = Set(Array(33...126) + Array(161...172) + Array(174...255))120        var extra = 0121        self.byteAlphabet = (0...255).map { byte in122            if visible.contains(byte) { return String(UnicodeScalar(byte)!) }123            defer { extra += 1 }124            return String(UnicodeScalar(256 + extra)!)125        }126    }127 128    private func bpe(_ token: String) throws -> [Int32] {129        if let existing = cache[token] { return existing }130        // ignore_merges gives a vocabulary hit precedence over applying merges.131        if let id = vocab[token] { return [id] }132        // HF BPE omits absent initial symbols if UNK and byte fallback are both133        // unset. This vocabulary lacks some controls, including mapped NUL.134        var symbols = token.unicodeScalars.map(String.init).filter { vocab[$0] != nil }135        while symbols.count > 1 {136            var bestRank = Int.max137            var bestIndex: Int?138            for index in 0..<(symbols.count - 1) {139                if let rank = ranks[symbols[index] + "\0" + symbols[index + 1]], rank < bestRank {140                    bestRank = rank141                    bestIndex = index142                }143            }144            guard let index = bestIndex else { break }145            symbols[index] += symbols[index + 1]146            symbols.remove(at: index + 1)147        }148        let ids = try symbols.map { symbol -> Int32 in149            guard let id = vocab[symbol] else { throw TokenizerError.missingToken(symbol) }150            return id151        }152        if cache.count < 32768 { cache[token] = ids }153        return ids154    }155 156    private func ordinary(_ text: String) throws -> [Int32] {157        let nsText = text as NSString158        let full = NSRange(location: 0, length: nsText.length)159        var ids: [Int32] = []160        var cursor = 0161        func appendPiece(_ range: NSRange) throws {162            guard range.length > 0 else { return }163            let raw = nsText.substring(with: range)164            let encoded = raw.utf8.map { byteAlphabet[Int($0)] }.joined()165            ids.append(contentsOf: try bpe(encoded))166        }167        for match in pattern.matches(in: text, range: full) {168            if match.range.location > cursor {169                try appendPiece(NSRange(location: cursor, length: match.range.location - cursor))170            }171            try appendPiece(match.range)172            cursor = NSMaxRange(match.range)173        }174        if cursor < nsText.length {175            try appendPiece(NSRange(location: cursor, length: nsText.length - cursor))176        }177        return ids178    }179 180    func encodeBody(text: String) throws -> [Int32] {181        let nsText = text as NSString182        let full = NSRange(location: 0, length: nsText.length)183        var ids: [Int32] = []184        var cursor = 0185        for match in addedPattern.matches(in: text, range: full) {186            let content = nsText.substring(with: match.range)187            guard let token = added[content] else { throw TokenizerError.missingToken(content) }188            if token.single_word {189                let before = nsText.substring(to: match.range.location)190                let after = nsText.substring(from: NSMaxRange(match.range))191                let leftWord = wordEnd.firstMatch(in: before, range: NSRange(location: 0, length: (before as NSString).length)) != nil192                let rightWord = wordStart.firstMatch(in: after, range: NSRange(location: 0, length: (after as NSString).length)) != nil193                if leftWord || rightWord { continue }194            }195            ids.append(contentsOf: try ordinary(nsText.substring(with: NSRange(location: cursor, length: match.range.location - cursor))))196            ids.append(token.id)197            cursor = NSMaxRange(match.range)198        }199        ids.append(contentsOf: try ordinary(nsText.substring(from: cursor)))200        return ids201    }202 203    func encode(text: String, sequenceLength: Int) throws -> (inputIDs: [Int32], attentionMask: [Int32]) {204        guard sequenceLength >= 2 else { throw TokenizerError.unsupported("sequenceLength must be >=2") }205        let body = try encodeBody(text: text)206        var ids = [clsID] + Array(body.prefix(sequenceLength - 2)) + [sepID]207        var mask = Array(repeating: Int32(1), count: ids.count)208        let padding = sequenceLength - ids.count209        ids.append(contentsOf: repeatElement(padID, count: padding))210        mask.append(contentsOf: repeatElement(Int32(0), count: padding))211        return (ids, mask)212    }213}214 215#if GRANITE_TOKENIZER_CLI216@main217enum GraniteTokenizerCLI {218    struct Request: Decodable { let id: String; let text: String; let sequence_length: Int }219    struct Response: Encodable { let id: String; let input_ids: [Int32]; let attention_mask: [Int32] }220    static func main() throws {221        guard CommandLine.arguments.count == 4 else {222            throw GraniteTokenizer.TokenizerError.unsupported("usage: tokenizer tokenizer.json tokenizer_config.json requests.json")223        }224        let tokenizer = try GraniteTokenizer(225            tokenizerURL: URL(fileURLWithPath: CommandLine.arguments[1]),226            tokenizerConfigURL: URL(fileURLWithPath: CommandLine.arguments[2]))227        let requests = try JSONDecoder().decode([Request].self, from: Data(contentsOf: URL(fileURLWithPath: CommandLine.arguments[3])))228        let responses = try requests.map { request -> Response in229            let result = try tokenizer.encode(text: request.text, sequenceLength: request.sequence_length)230            return Response(id: request.id, input_ids: result.inputIDs, attention_mask: result.attentionMask)231        }232        FileHandle.standardOutput.write(try JSONEncoder().encode(responses))233    }234}235#endif236