File size: 2,122 Bytes
357ae2c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | import Foundation
/// Old→new token-id remap for the EN+ZH pruned models.
///
/// The `*-en_zh-*` Core ML models were built on a sliced word-embedding matrix
/// (only ASCII/CJK tokens + specials kept). The app must therefore translate the
/// full-vocabulary XLM-R ids produced by the tokenizer into the pruned id space
/// before feeding the model; dropped tokens (`-1`) become `<unk>`.
///
/// Special ids 0/1/2/3 (`<s>`/`<pad>`/`</s>`/`<unk>`) map to themselves, so the
/// same table can be applied uniformly to every id in a window, including the
/// padding and sentinel tokens.
struct VocabRemap {
private let table: [Int32] // table[oldId] = newId, or -1 if dropped
let unkNewId: Int32
/// Layout (little-endian): magic "WTPR", u32 count, then `count` × i32.
init(url: URL) throws {
let data = try Data(contentsOf: url, options: .mappedIfSafe)
var table = [Int32]()
try data.withUnsafeBytes { (raw: UnsafeRawBufferPointer) in
guard raw.count >= 8 else { throw KitError.corruptResource("en_zh_remap.bin too small") }
let magic = raw.loadUnaligned(fromByteOffset: 0, as: UInt32.self)
guard magic == 0x52505457 else { throw KitError.corruptResource("bad remap magic") } // "WTPR"
let count = Int(UInt32(littleEndian: raw.loadUnaligned(fromByteOffset: 4, as: UInt32.self)))
guard raw.count >= 8 + count * 4 else { throw KitError.corruptResource("remap truncated") }
table = [Int32](repeating: 0, count: count)
for k in 0..<count {
let v = raw.loadUnaligned(fromByteOffset: 8 + k * 4, as: Int32.self)
table[k] = Int32(littleEndian: v)
}
}
self.table = table
// <unk> old id is 3; its mapped value is the pruned unk id.
self.unkNewId = (3 < table.count && table[3] >= 0) ? table[3] : 3
}
@inline(__always)
func map(_ oldId: Int32) -> Int32 {
guard oldId >= 0 && Int(oldId) < table.count else { return unkNewId }
let v = table[Int(oldId)]
return v >= 0 ? v : unkNewId
}
}
|