| import Foundation |
| import SpeechCore |
|
|
| |
| |
| |
| package final class AssetStore { |
| package struct Manifest: Codable { |
| package struct PromptIDs: Codable { |
| package let user: Int |
| package let bos_audio: Int |
| package let audio: Int |
| package let eos_audio: Int |
| package let text: [Int] |
| package let assistant: Int |
| } |
| package struct Tokens: Codable { |
| package let audio_token_id: Int |
| package let pad_token_id: Int |
| package let eos_token_ids: [Int] |
| package let extra_block_token_ids: [Int] |
| } |
| package let prompt_ids: PromptIDs |
| package let tokens: Tokens |
| } |
|
|
| package let manifest: Manifest |
| package let hannWindow: [Float] |
| package let melFilters: [Float] |
| package let projNormW: [Float] |
| package let projNormB: [Float] |
| package let projLinW: [Float] |
| package let projLinB: [Float] |
|
|
| private let embeddingData: Data |
| private let vocabBytes: Data |
| private let vocabOffsets: [UInt32] |
| package let embeddingDim = 512 |
|
|
| package init(bundle: AssetBundle) throws { |
| func data(_ name: String) throws -> Data { |
| try Data(contentsOf: bundle.url(name), options: .mappedIfSafe) |
| } |
| func floats(_ name: String) throws -> [Float] { |
| try data(name).withUnsafeBytes { Array($0.bindMemory(to: Float.self)) } |
| } |
| do { |
| manifest = try JSONDecoder().decode(Manifest.self, from: data("asr_manifest")) |
| } catch let e as SpeechError { |
| throw e |
| } catch { |
| throw SpeechError.assetMissing("asr_manifest (decode: \(error.localizedDescription))") |
| } |
| hannWindow = try floats("hann_window") |
| melFilters = try floats("mel_filters") |
| projNormW = try floats("projector_norm_w") |
| projNormB = try floats("projector_norm_b") |
| projLinW = try floats("projector_lin_w") |
| projLinB = try floats("projector_lin_b") |
| embeddingData = try data("token_embedding") |
| vocabBytes = try data("vocab_bytes") |
| vocabOffsets = try data("vocab_offsets").withUnsafeBytes { |
| Array($0.bindMemory(to: UInt32.self)) |
| } |
| } |
|
|
| |
| package func embedding(for tokenID: Int) -> [Float] { |
| let rowBytes = embeddingDim * MemoryLayout<Float16>.size |
| let start = tokenID * rowBytes |
| return embeddingData.subdata(in: start..<(start + rowBytes)).withUnsafeBytes { |
| $0.bindMemory(to: Float16.self).map { Float($0) } |
| } |
| } |
|
|
| package func decode(_ tokenIDs: [Int]) -> String { |
| var bytes = Data() |
| for id in tokenIDs where id >= 0 && id + 1 < vocabOffsets.count { |
| let a = Int(vocabOffsets[id]), b = Int(vocabOffsets[id + 1]) |
| bytes.append(vocabBytes.subdata(in: a..<b)) |
| } |
| let raw = String(decoding: bytes, as: UTF8.self) |
| return raw.trimmingCharacters(in: .whitespacesAndNewlines) |
| .replacingOccurrences(of: #"\s+"#, with: " ", options: .regularExpression) |
| } |
| } |
|
|
