Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Package.swift
Original file line number Diff line number Diff line change
Expand Up @@ -266,6 +266,11 @@ let package = Package(
.linkedLibrary("c++")
]
),
.testTarget(
Comment thread
stikves marked this conversation as resolved.
name: "SpeechTests",
dependencies: ["CoreAISpeech"],
path: "swift/Tests/SpeechTests"
),
],
swiftLanguageModes: [.v6],
cxxLanguageStandard: .cxx17
Expand Down
15 changes: 14 additions & 1 deletion swift/Sources/CoreAISpeech/SpeechBundle.swift
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,20 @@ public struct GenerationConfig: Sendable {
init(from url: URL) throws {
let data = try Data(contentsOf: url)
let json = try JSONSerialization.jsonObject(with: data) as? [String: Any] ?? [:]
forcedPrefix = (json["forced_decoder_ids"] as? [Int]).map { $0.map { Int32($0) } } ?? Self.whisper.forcedPrefix
if let rawPairs = json["forced_decoder_ids"] as? [[Any]] {
// HF format: [[position, token_id], ...] where token_id can be null
var tokens: [(pos: Int, id: Int32)] = []
for pair in rawPairs {
guard pair.count >= 2,
let pos = pair[0] as? Int,
let id = pair[1] as? Int
else { continue }
tokens.append((pos: pos, id: Int32(id)))
}
forcedPrefix = tokens.sorted { $0.pos < $1.pos }.map(\.id)
} else {
forcedPrefix = Self.whisper.forcedPrefix
}
eotToken = (json["eos_token_id"] as? Int).map { Int32($0) } ?? Self.whisper.eotToken
maxDecodeSteps = (json["max_new_tokens"] as? Int) ?? Self.whisper.maxDecodeSteps
tokenizerName = json["tokenizer_name"] as? String ?? Self.whisper.tokenizerName
Expand Down
84 changes: 84 additions & 0 deletions swift/Tests/SpeechTests/GenerationConfigTests.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
// Copyright 2026 Apple Inc.
//
// Use of this source code is governed by a BSD-3-clause license that can
// be found in the LICENSE file or at https://opensource.org/licenses/BSD-3-Clause

import Foundation
import Testing

@testable import CoreAISpeech

@Suite("GenerationConfig")
struct GenerationConfigTests {
@Test("Parses forced_decoder_ids from HuggingFace format")
func parseForcedDecoderIds() throws {
let json: [String: Any] = [
"forced_decoder_ids": [[1, 50258], [2, 50259], [3, 50360], [4, 50364]],
"eos_token_id": 50257,
"max_new_tokens": 100,
]
let data = try JSONSerialization.data(withJSONObject: json)
let url = FileManager.default.temporaryDirectory.appendingPathComponent("gen_config.json")
try data.write(to: url)
defer { try? FileManager.default.removeItem(at: url) }

let config = try GenerationConfig(from: url)
#expect(config.forcedPrefix == [50258, 50259, 50360, 50364])
#expect(config.eotToken == 50257)
#expect(config.maxDecodeSteps == 100)
}

@Test("Falls back to Whisper defaults when forced_decoder_ids is missing")
func fallbackOnMissingField() throws {
let json: [String: Any] = ["eos_token_id": 50257]
let data = try JSONSerialization.data(withJSONObject: json)
let url = FileManager.default.temporaryDirectory.appendingPathComponent("gen_config2.json")
try data.write(to: url)
defer { try? FileManager.default.removeItem(at: url) }

let config = try GenerationConfig(from: url)
#expect(config.forcedPrefix == GenerationConfig.whisper.forcedPrefix)
}

@Test("Falls back to Whisper defaults when forced_decoder_ids has wrong format")
func fallbackOnWrongFormat() throws {
let json: [String: Any] = [
"forced_decoder_ids": [50258, 50259, 50360] // flat array, not pairs
]
let data = try JSONSerialization.data(withJSONObject: json)
let url = FileManager.default.temporaryDirectory.appendingPathComponent("gen_config3.json")
try data.write(to: url)
defer { try? FileManager.default.removeItem(at: url) }

let config = try GenerationConfig(from: url)
#expect(config.forcedPrefix == GenerationConfig.whisper.forcedPrefix)
}

@Test("Skips null token IDs in forced_decoder_ids pairs")
func skipsNullTokenIds() throws {
let json: [String: Any] = [
"forced_decoder_ids": [[1, NSNull()], [2, 50360], [3, 50364]]
]
let data = try JSONSerialization.data(withJSONObject: json)
let url = FileManager.default.temporaryDirectory.appendingPathComponent("gen_config4.json")
try data.write(to: url)
defer { try? FileManager.default.removeItem(at: url) }

let config = try GenerationConfig(from: url)
#expect(config.forcedPrefix == [50360, 50364])
}

@Test("Handles out-of-order positions")
func outOfOrderPositions() throws {
let json: [String: Any] = [
"forced_decoder_ids": [[3, 50360], [1, 50258], [2, 50259]]
]
let data = try JSONSerialization.data(withJSONObject: json)
let url = FileManager.default.temporaryDirectory.appendingPathComponent("gen_config5.json")
try data.write(to: url)
defer { try? FileManager.default.removeItem(at: url) }

let config = try GenerationConfig(from: url)
#expect(config.forcedPrefix == [50258, 50259, 50360])
}
}