Sources/WtpsplitKit/Tokenizer/UnigramTokenizer.swift

import Foundation

/// XLM-RoBERTa tokenizer (SentencePiece **Unigram**) reimplemented in Swift,
/// matching the pipeline `scripts/run_segmentation.py` relies on:
///
///   normalize (NFKC) → WhitespaceSplit + Metaspace(▁) → Unigram Viterbi → `<s> … </s>`
///
/// The crucial extra it produces beyond ids is, for every content token, the
/// **original-scalar index of the token's last character** (`charEnds`). The SaT
/// head emits one boundary logit per token, and that logit is assigned to the
/// last character of the token's span (see `token_to_char_probs` in the Python
/// reference). Without offsets the logits cannot be placed back onto characters.
///
/// All indices are Unicode-scalar (code-point) based to match Python `str`.
final class UnigramTokenizer {

    // XLM-R special-token ids (fixed by the vocabulary).
    static let bosId: Int32 = 0   // <s>
    static let padId: Int32 = 1   // <pad>
    static let eosId: Int32 = 2   // </s>

    /// SentencePiece metaspace marker (U+2581 "▁").
    private static let metaspace: Unicode.Scalar = "\u{2581}"

    /// SentencePiece unknown-token penalty added to `minScore`.
    private static let unkPenalty: Float = 10.0

    private let pieceToId: [String: Int32]
    private let scores: [Float]            // indexed by token id
    let unkId: Int32
    private let unkScore: Float
    private let maxPieceScalars: Int

    /// One encoded token stream (specials excluded; callers add `<s>`/`</s>`).
    struct Encoding {
        /// Content token ids, in order (no `<s>`/`</s>`).
        var ids: [Int32]
        /// `charEnds[k]` = index, in the *original* scalar array, of the last
        /// character covered by token `k`. Used to scatter logits onto chars.
        var charEnds: [Int]
    }

    /// Load the tokenizer from the vocabulary bundled with the package.
    static func bundled() throws -> UnigramTokenizer {
        guard let url = Bundle.module.url(forResource: "xlmr_unigram", withExtension: "bin") else {
            throw KitError.missingResource("xlmr_unigram.bin")
        }
        return try UnigramTokenizer(vocabURL: url)
    }

    /// Load the packed Unigram vocabulary (`xlmr_unigram.bin`).
    ///
    /// Layout (little-endian): magic "WTPV", u32 version, u32 count, u32 unkId,
    /// then `count` records of `[u16 byteLen][f32 score][byteLen UTF-8 bytes]`.
    init(vocabURL: URL) throws {
        let data = try Data(contentsOf: vocabURL, options: .mappedIfSafe)
        var pieceToId = [String: Int32]()
        var scores = [Float]()
        var unkId: Int32 = 3
        var maxScalars = 1

        try data.withUnsafeBytes { (raw: UnsafeRawBufferPointer) in
            guard raw.count >= 16 else { throw KitError.corruptResource("xlmr_unigram.bin too small") }
            var off = 0
            func u32() -> UInt32 {
                let v = raw.loadUnaligned(fromByteOffset: off, as: UInt32.self); off += 4
                return UInt32(littleEndian: v)
            }
            func u16() -> UInt16 {
                let v = raw.loadUnaligned(fromByteOffset: off, as: UInt16.self); off += 2
                return UInt16(littleEndian: v)
            }
            func f32() -> Float {
                let bits = raw.loadUnaligned(fromByteOffset: off, as: UInt32.self); off += 4
                return Float(bitPattern: UInt32(littleEndian: bits))
            }

            let magic = raw.loadUnaligned(fromByteOffset: 0, as: UInt32.self); off = 4
            // "WTPV" little-endian
            guard magic == 0x56505457 else { throw KitError.corruptResource("bad vocab magic") }
            _ = u32()                         // version
            let count = Int(u32())
            unkId = Int32(u32())

            pieceToId.reserveCapacity(count)
            scores = [Float](repeating: 0, count: count)
            let base = raw.baseAddress!

            for id in 0..<count {
                let len = Int(u16())
                let score = f32()
                let piece = String(decoding: UnsafeRawBufferPointer(start: base + off, count: len),
                                   as: UTF8.self)
                off += len
                scores[id] = score
                pieceToId[piece] = Int32(id)
                let n = piece.unicodeScalars.count
                if n > maxScalars { maxScalars = n }
            }
        }

        self.pieceToId = pieceToId
        self.scores = scores
        self.unkId = unkId
        self.maxPieceScalars = maxScalars
        let minScore = scores.min() ?? 0
        self.unkScore = minScore - Self.unkPenalty
    }

    /// Tokenize the original text (already split into scalars) into content tokens.
    func encode(scalars original: [Unicode.Scalar]) -> Encoding {
        var enc = Encoding(ids: [], charEnds: [])
        if original.isEmpty { return enc }

        // 1) Normalize (NFKC, per scalar) keeping a normalized→original index map.
        let (norm, origIdx) = normalize(original)
        let n = norm.count
        if n == 0 { return enc }

        // 2) WhitespaceSplit: iterate maximal non-space runs as words.
        var i = 0
        while i < n {
            if CharClasses.isSpace(norm[i]) { i += 1; continue }
            var j = i
            while j < n && !CharClasses.isSpace(norm[j]) { j += 1 }
            encodeWord(norm: norm, origIdx: origIdx, wordStart: i, wordEnd: j, into: &enc)
            i = j
        }
        return enc
    }

    // MARK: - Word-level Unigram Viterbi

    private func encodeWord(norm: [Unicode.Scalar], origIdx: [Int],
                            wordStart: Int, wordEnd: Int, into enc: inout Encoding) {
        // 3) Metaspace(add_prefix_space): S = ▁ + word.
        let wordLen = wordEnd - wordStart
        let L = wordLen + 1
        var s = [Unicode.Scalar](repeating: Self.metaspace, count: L)
        for k in 0..<wordLen { s[k + 1] = norm[wordStart + k] }

        // Viterbi over S maximizing summed Unigram log-scores.
        let negInf = -Float.greatestFiniteMagnitude
        var best = [Float](repeating: negInf, count: L + 1)
        var backPrev = [Int](repeating: -1, count: L + 1)
        var backId = [Int32](repeating: unkId, count: L + 1)
        best[0] = 0

        for pos in 0..<L {
            if best[pos] == negInf { continue }
            let maxLen = min(maxPieceScalars, L - pos)
            var view = String.UnicodeScalarView()
            view.reserveCapacity(maxLen)
            if maxLen >= 1 {
                for len in 1...maxLen {
                    view.append(s[pos + len - 1])
                    if let id = pieceToId[String(view)] {
                        let cand = best[pos] + scores[Int(id)]
                        if cand > best[pos + len] {
                            best[pos + len] = cand
                            backPrev[pos + len] = pos
                            backId[pos + len] = id
                        }
                    }
                }
            }
            // Single-character <unk> fallback (guarantees a complete path).
            let cand = best[pos] + unkScore
            if cand > best[pos + 1] {
                best[pos + 1] = cand
                backPrev[pos + 1] = pos
                backId[pos + 1] = unkId
            }
        }

        // Backtrack to recover pieces in forward order.
        var pieces: [(end: Int, id: Int32)] = []
        var pos = L
        while pos > 0 {
            pieces.append((pos, backId[pos]))
            pos = backPrev[pos]
        }
        pieces.reverse()

        // 4) Emit ids + map each token's last char back to an original index.
        for pc in pieces {
            enc.ids.append(pc.id)
            if pc.end >= 2 {
                // S index (pc.end-1) ≥ 1 → real word scalar.
                let normIdx = wordStart + pc.end - 2
                enc.charEnds.append(origIdx[normIdx])
            } else {
                // Lone "▁" piece: HF Metaspace maps it onto the word's first
                // character (the boundary sits at the word start).
                enc.charEnds.append(origIdx[wordStart])
            }
        }
    }

    // MARK: - Normalization

    /// Per-scalar NFKC (Foundation `precomposedStringWithCompatibilityMapping`),
    /// matching XLM-R's SentencePiece nmt_nfkc closely for Latin/CJK text. Returns
    /// the normalized scalars plus, for each, the index of the original scalar it
    /// came from — so split positions land on the *original* string.
    private func normalize(_ original: [Unicode.Scalar]) -> (norm: [Unicode.Scalar], origIdx: [Int]) {
        var norm = [Unicode.Scalar]()
        var origIdx = [Int]()
        norm.reserveCapacity(original.count)
        origIdx.reserveCapacity(original.count)
        for (i, sc) in original.enumerated() {
            let mapped = String(sc).precomposedStringWithCompatibilityMapping
            for ns in mapped.unicodeScalars {
                norm.append(ns)
                origIdx.append(i)
            }
        }
        return (norm, origIdx)
    }
}