Add MultilingualConformanceTests for byte-identical Python parity (#352)

.gitignore

@@ -12,3 +12,6 @@ DerivedData/
 .idea
 .index-build
 *.out
+
+# Local venv for Tools/generate_tokenizer_baselines.py
+.venv*/

Tests/TokenizersTests/MultilingualConformanceTests.swift

@@ -0,0 +1,330 @@
+//
+//  MultilingualConformanceTests.swift
+//
+//  Byte-identical conformance tests: the Swift port's `tokenizer.encode(text:)`
+//  output must match the canonical HuggingFace Python `transformers` reference
+//  for every input in `Resources/MultilingualConformance/inputs.json` across
+//  the kernels in `BaselineKernel.all`.
+//
+//  Inputs are categorised and stable-keyed so a divergence message points
+//  reviewers at the exact axis that broke (japanese-voiced-kana, emoji-keycap,
+//  thai-combining-marks, …). Baselines are regenerated by
+//  `Tools/generate_tokenizer_baselines.py` against the transformers version
+//  pinned in `Tools/requirements.txt`.
+//
+//  `expectedDivergences` lets the test target ship green while bug fixes are
+//  in flight: each entry pairs a (model, input id) with the PR that will land
+//  the fix. An unexpected divergence is a hard failure (regression catch);
+//  an unexpected *match* prints a cleanup hint inviting removal of the entry
+//  but doesn't fail (so a freshly merged fix doesn't break CI on this file).
+//
+//  Conformance design conceptually anchored in @apocryphx's
+//  ObjCTokenizer port (https://github.com/apocryphx/ObjCTokenizer) and the
+//  `expectedDivergences` + decoded-fields ideas from @john-rocky's closed
+//  #357. See issue #352 for the multilingual-divergence catalogue this
+//  corpus exercises.
+
+import Foundation
+import Testing
+
+@testable import Hub
+@testable import Models
+@testable import Tokenizers
+
+// MARK: - Resource model
+
+private struct CorpusEntry: Decodable, Sendable {
+    let id: String
+    let category: String
+    let text: String
+}
+
+private struct BaselineEntry: Decodable, Sendable {
+    let id: String
+    let inputIds: [Int]
+    let tokens: [String]
+    let decodedWithSpecial: String
+    let decodedSkipSpecial: String
+
+    enum CodingKeys: String, CodingKey {
+        case id
+        case inputIds = "input_ids"
+        case tokens
+        case decodedWithSpecial = "decoded_with_special"
+        case decodedSkipSpecial = "decoded_skip_special"
+    }
+}
+
+private struct BaselineMetadata: Decodable, Sendable {
+    let modelId: String
+    let transformersVersion: String
+    let generatedAt: String
+    let inputCount: Int
+
+    enum CodingKeys: String, CodingKey {
+        case modelId = "model_id"
+        case transformersVersion = "transformers_version"
+        case generatedAt = "generated_at"
+        case inputCount = "input_count"
+    }
+}
+
+private struct BaselineFile: Decodable, Sendable {
+    let metadata: BaselineMetadata
+    let entries: [BaselineEntry]
+}
+
+// MARK: - Kernel matrix
+
+private struct BaselineKernel: Sendable, CustomStringConvertible {
+    /// Filename slug under `Resources/MultilingualConformance/baselines/<slug>_multilingual.json`.
+    let slug: String
+    /// `AutoTokenizer.from(pretrained:)` argument.
+    let modelId: String
+
+    var description: String { modelId }
+
+    static let all: [BaselineKernel] = [
+        BaselineKernel(slug: "bge_small",    modelId: "BAAI/bge-small-en-v1.5"),
+        BaselineKernel(slug: "t5_small",    modelId: "google-t5/t5-small"),
+        BaselineKernel(slug: "gpt2",        modelId: "openai-community/gpt2"),
+        BaselineKernel(slug: "roberta_base", modelId: "FacebookAI/roberta-base"),
+        BaselineKernel(slug: "qwen2_5",     modelId: "Qwen/Qwen2.5-0.5B"),
+        BaselineKernel(slug: "tinyllama",   modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0"),
+    ]
+}
+
+// MARK: - Divergences known to be in flight
+
+/// (modelId, inputId) pairs whose encode output is known to diverge from the
+/// Python reference today, with a free-form note documenting the surface so
+/// follow-up triage has a starting point. Cleanup-hint pattern inspired by
+/// @john-rocky's closed #357: an unexpected match prints a hint inviting
+/// entry removal, an unexpected divergence is a hard test failure.
+private struct ExpectedDivergence: Sendable, Hashable {
+    let modelId: String
+    let inputId: String
+    let note: String
+}
+
+private let expectedDivergences: Set<ExpectedDivergence> = [
+    //
+    // Two new bug clusters this corpus surfaces that aren't addressed by the
+    // initial fix wave (#354 / #355 / #356, all merged). Worth filing as
+    // separate follow-up issues under #352.
+    //
+
+    // SentencePiece-BPE leading-whitespace runs collapse to single `▁` tokens
+    // instead of producing a single multi-space vocab entry (e.g. `▁▁▁▁`).
+    // Suggests the Metaspace pre-tokenizer or BPE merge step isn't recognising
+    // `▁▁▁▁` (id 268 in TinyLlama vocab) as a vocab-eligible merge target.
+    .init(modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0", inputId: "code-python-if",          note: "Metaspace leading-whitespace runs"),
+    .init(modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0", inputId: "code-python-return",      note: "Metaspace leading-whitespace runs"),
+    .init(modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0", inputId: "code-python-recurse",     note: "Metaspace leading-whitespace runs"),
+    .init(modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0", inputId: "whitespace-runs",         note: "Metaspace leading-whitespace runs"),
+    .init(modelId: "TinyLlama/TinyLlama-1.1B-Chat-v1.0", inputId: "whitespace-trailing-tabs", note: "Metaspace leading-whitespace runs"),
+
+    // Qwen2.5 byte-level BPE picks a different merge ordering on Thai
+    // (and Thai-inside-multiscript) than HF Python. Byte-level encoding
+    // means there are no combining-mark traps; this is a merge-priority
+    // ordering issue in the BPE algorithm itself. Worth tracing once
+    // #355's merge-loop changes have settled.
+    .init(modelId: "Qwen/Qwen2.5-0.5B", inputId: "thai-combining-marks-greeting", note: "byte-level BPE merge-ordering on Thai"),
+    .init(modelId: "Qwen/Qwen2.5-0.5B", inputId: "thai-combining-marks-prose",    note: "byte-level BPE merge-ordering on Thai"),
+    .init(modelId: "Qwen/Qwen2.5-0.5B", inputId: "multiscript-greetings",         note: "byte-level BPE merge-ordering on Thai"),
+]
+
+private func divergenceExpected(model: String, input: String) -> ExpectedDivergence? {
+    expectedDivergences.first { $0.modelId == model && $0.inputId == input }
+}
+
+// MARK: - Resource loading
+
+private enum ConformanceError: Error, CustomStringConvertible {
+    case missingResource(String)
+    case decodeError(String, Error)
+    case unsupportedTokenizer(String)
+
+    var description: String {
+        switch self {
+        case .missingResource(let name): "missing test resource: \(name)"
+        case .decodeError(let name, let err): "decode error in \(name): \(err)"
+        case .unsupportedTokenizer(let id): "tokenizer for \(id) was not a PreTrainedTokenizer"
+        }
+    }
+}
+
+// SwiftPM's `.process("Resources")` flattens subdirectory structure into the
+// test bundle root, so `subdirectory:` lookups don't apply here. Filenames
+// (`inputs.json`, `<slug>_multilingual.json`) are unique within the test
+// resources, so flat lookup is unambiguous.
+private func loadCorpus() throws -> [CorpusEntry] {
+    guard let url = Bundle.module.url(forResource: "inputs", withExtension: "json") else {
+        throw ConformanceError.missingResource("inputs.json")
+    }
+    let data = try Data(contentsOf: url)
+    do { return try JSONDecoder().decode([CorpusEntry].self, from: data) }
+    catch { throw ConformanceError.decodeError("inputs.json", error) }
+}
+
+private func loadBaseline(slug: String) throws -> BaselineFile {
+    let resource = "\(slug)_multilingual"
+    guard let url = Bundle.module.url(forResource: resource, withExtension: "json") else {
+        throw ConformanceError.missingResource("\(resource).json")
+    }
+    let data = try Data(contentsOf: url)
+    do { return try JSONDecoder().decode(BaselineFile.self, from: data) }
+    catch { throw ConformanceError.decodeError(resource, error) }
+}
+
+// MARK: - Failure diagnostics
+
+/// Format a windowed diff around the first divergence point: `expected_window`
+/// + `got_window` decoded to readable token strings, with the divergence
+/// position underlined. Easier to triage than two long id arrays.
+private func divergenceReport(
+    inputId: String,
+    category: String,
+    expectedIds: [Int],
+    expectedTokens: [String],
+    gotIds: [Int],
+    gotTokens: [String]
+) -> String {
+    let commonLen = min(expectedIds.count, gotIds.count)
+    var divIdx = 0
+    while divIdx < commonLen, expectedIds[divIdx] == gotIds[divIdx] { divIdx += 1 }
+    let windowLo = max(0, divIdx - 3)
+    let windowExpectedHi = min(expectedIds.count, divIdx + 5)
+    let windowGotHi = min(gotIds.count, divIdx + 5)
+
+    func annotate(_ ids: [Int], _ tokens: [String], hi: Int) -> String {
+        var parts: [String] = []
+        for i in windowLo..<hi {
+            let tok = i < tokens.count ? tokens[i] : "?"
+            parts.append("[\(ids[i])]\(tok)")
+        }
+        return parts.joined(separator: " ")
+    }
+
+    return """
+    multilingual conformance divergence
+      input id:    \(inputId)
+      category:    \(category)
+      divergence:  index \(divIdx) (expected len=\(expectedIds.count), got len=\(gotIds.count))
+      expected:    \(annotate(expectedIds, expectedTokens, hi: windowExpectedHi))
+      got:         \(annotate(gotIds, gotTokens, hi: windowGotHi))
+    """
+}
+
+// MARK: - Tests
+
+@Suite("Multilingual Conformance")
+struct MultilingualConformanceTests {
+    @Test("Corpus inputs.json is well-formed (unique ids, non-empty fields)")
+    func corpusIsWellFormed() throws {
+        let corpus = try loadCorpus()
+        #expect(!corpus.isEmpty, "corpus should not be empty")
+        var seen = Set<String>()
+        for entry in corpus {
+            #expect(!entry.id.isEmpty, "empty id at category=\(entry.category)")
+            #expect(!entry.category.isEmpty, "empty category at id=\(entry.id)")
+            // text may legitimately contain trailing whitespace etc.;
+            // empty string isn't a valid encode target though.
+            #expect(!entry.text.isEmpty, "empty text at id=\(entry.id)")
+            #expect(!seen.contains(entry.id), "duplicate id: \(entry.id)")
+            seen.insert(entry.id)
+        }
+    }
+
+    @Test("Baselines cover the corpus exactly", arguments: BaselineKernel.all)
+    fileprivate func baselinesCoverCorpus(kernel: BaselineKernel) throws {
+        let corpus = try loadCorpus()
+        let baseline = try loadBaseline(slug: kernel.slug)
+        let corpusIds = Set(corpus.map(\.id))
+        let baselineIds = Set(baseline.entries.map(\.id))
+
+        let missing = corpusIds.subtracting(baselineIds)
+        let extra   = baselineIds.subtracting(corpusIds)
+        #expect(
+            missing.isEmpty,
+            "baseline for \(kernel.modelId) is missing entries: \(missing.sorted())"
+        )
+        #expect(
+            extra.isEmpty,
+            "baseline for \(kernel.modelId) has stale entries (corpus shrunk?): \(extra.sorted())"
+        )
+        #expect(
+            baseline.metadata.inputCount == baseline.entries.count,
+            "baseline metadata input_count mismatch for \(kernel.modelId): metadata=\(baseline.metadata.inputCount) entries=\(baseline.entries.count)"
+        )
+    }
+
+    @Test("Byte-identical token ids vs HF Python", arguments: BaselineKernel.all)
+    fileprivate func byteIdenticalTokenIds(kernel: BaselineKernel) async throws {
+        let corpus = try loadCorpus()
+        let baseline = try loadBaseline(slug: kernel.slug)
+        let entriesById = Dictionary(uniqueKeysWithValues: baseline.entries.map { ($0.id, $0) })
+
+        let tokenizerOpt = try await AutoTokenizer.from(pretrained: kernel.modelId) as? PreTrainedTokenizer
+        guard let tokenizer = tokenizerOpt else {
+            throw ConformanceError.unsupportedTokenizer(kernel.modelId)
+        }
+
+        var unexpectedDivergences: [String] = []
+        var unexpectedMatches: [ExpectedDivergence] = []
+
+        for input in corpus {
+            guard let expected = entriesById[input.id] else { continue }
+            let got = tokenizer.encode(text: input.text)
+            let knownDivergence = divergenceExpected(model: kernel.modelId, input: input.id)
+
+            if got == expected.inputIds {
+                if let exp = knownDivergence {
+                    // Test stays green but the table needs cleanup.
+                    unexpectedMatches.append(exp)
+                }
+            } else {
+                if knownDivergence != nil {
+                    // Listed in expectedDivergences — this is the in-flight
+                    // state, not a regression.
+                    continue
+                }
+                let gotTokens = got.map { tokenizer.convertIdToToken($0) ?? "?" }
+                unexpectedDivergences.append(
+                    divergenceReport(
+                        inputId: input.id,
+                        category: input.category,
+                        expectedIds: expected.inputIds,
+                        expectedTokens: expected.tokens,
+                        gotIds: got,
+                        gotTokens: gotTokens
+                    )
+                )
+            }
+        }
+
+        // Cleanup hint — does NOT fail the test. A freshly merged improvement
+        // shouldn't break CI on this file; the hint just tells reviewers an
+        // expectedDivergences entry can be removed.
+        for match in unexpectedMatches {
+            print("""
+            [\(kernel.modelId)] expectedDivergences entry no longer applies:
+              input id:  \(match.inputId)
+              note:      \(match.note)
+              hint:      remove this entry from expectedDivergences in MultilingualConformanceTests.swift
+            """)
+        }
+
+        // Regression catch — IS a failure.
+        #expect(
+            unexpectedDivergences.isEmpty,
+            """
+            \(kernel.modelId): \(unexpectedDivergences.count) unexpected divergence(s) from HF Python reference.
+            If a divergence is being addressed in an open PR or is otherwise
+            known, add an ExpectedDivergence(modelId: …, inputId: …, note: …) entry.
+
+            \(unexpectedDivergences.joined(separator: "\n\n"))
+            """
+        )
+    }
+}