diff --git a/Cotabby.xcodeproj/project.pbxproj b/Cotabby.xcodeproj/project.pbxproj index 3acb62fb..844ed8b6 100644 --- a/Cotabby.xcodeproj/project.pbxproj +++ b/Cotabby.xcodeproj/project.pbxproj @@ -54,6 +54,7 @@ 279F017530A86AF62EB17918 /* EmojiSynonymCatalog.swift in Sources */ = {isa = PBXBuildFile; fileRef = 0846DE4E0293AF13890620D3 /* EmojiSynonymCatalog.swift */; }; 27D4F5CACADE171F142178B4 /* SettingsSidebarView.swift in Sources */ = {isa = PBXBuildFile; fileRef = BADB38D0160B47637572FC5E /* SettingsSidebarView.swift */; }; 286B7022E2A2774275004447 /* WelcomeTemplateStepView.swift in Sources */ = {isa = PBXBuildFile; fileRef = A9199B9CEAB320982CA333B8 /* WelcomeTemplateStepView.swift */; }; + 29EC35D67D9B4C3C50222619 /* ConstrainedBeamSearchTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */; }; 2C6159231472A849F15BD0AE /* ScreenFrameReader.swift in Sources */ = {isa = PBXBuildFile; fileRef = 5484C8A04B9C00CF79D589EB /* ScreenFrameReader.swift */; }; 2DF5A3826AAB99C279EBB8DE /* InputMonitor.swift in Sources */ = {isa = PBXBuildFile; fileRef = B81DD30EB657368AACE9625A /* InputMonitor.swift */; }; 2E3DEB7E89D0146274596F2E /* SettingsContainerView.swift in Sources */ = {isa = PBXBuildFile; fileRef = DB0CE9AB1286367BA2E82392 /* SettingsContainerView.swift */; }; @@ -200,6 +201,7 @@ C2C958D6E5F5FE1CCC414BCE /* SuggestionSubsystemContracts.swift in Sources */ = {isa = PBXBuildFile; fileRef = DEB16474A67CE1D210B944C9 /* SuggestionSubsystemContracts.swift */; }; C4C6734678797669055988E0 /* AppUpdateManager.swift in Sources */ = {isa = PBXBuildFile; fileRef = AD9573F3504CAE6891DF9B7D /* AppUpdateManager.swift */; }; C618C5595DA9C57C806A3E03 /* SettingsAttentionEvaluatorTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 2BC293F6125E2B14DCF05AD9 /* SettingsAttentionEvaluatorTests.swift */; }; + C638466A2F373BABA6F60A6D /* ConstrainedBeamSearch.swift in Sources */ = {isa = PBXBuildFile; fileRef = A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */; }; C71B594433F3B411CAE5DE7E /* FocusCapabilityResolverTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D4F6D5F94B238F7B4BE7C247 /* FocusCapabilityResolverTests.swift */; }; C9B815652CED38966C53A5E8 /* EmojiVariantResolverTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = EE8BB19D8EC9A75CD3458A6B /* EmojiVariantResolverTests.swift */; }; CA5B2D226FBAA5419E78F14F /* SuggestionSessionReconciler.swift in Sources */ = {isa = PBXBuildFile; fileRef = CE0AA0503128B0FC3951D700 /* SuggestionSessionReconciler.swift */; }; @@ -423,6 +425,7 @@ A52D0B550E00EF173A5D157E /* LlamaRuntimeManager.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LlamaRuntimeManager.swift; sourceTree = ""; }; A804F4DB6FD9BC8C27B2B65F /* LlamaRuntimeModels.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = LlamaRuntimeModels.swift; sourceTree = ""; }; A829F28F01FAE76CA7244BBC /* ModelFileValidatorTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelFileValidatorTests.swift; sourceTree = ""; }; + A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ConstrainedBeamSearch.swift; sourceTree = ""; }; A854CAFB1F557BC4CAED8819 /* VisualContextCoordinator.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = VisualContextCoordinator.swift; sourceTree = ""; }; A863F41C0C03D7B4AC5DC002 /* MarkerSelectionSynthesizer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = MarkerSelectionSynthesizer.swift; sourceTree = ""; }; A9199B9CEAB320982CA333B8 /* WelcomeTemplateStepView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WelcomeTemplateStepView.swift; sourceTree = ""; }; @@ -479,6 +482,7 @@ D5D6C2318E405AA717D1C256 /* WelcomePermissionStepView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = WelcomePermissionStepView.swift; sourceTree = ""; }; D814BBA41CF29E8DD9954651 /* OnboardingTemplateFeatureListTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = OnboardingTemplateFeatureListTests.swift; sourceTree = ""; }; D84D4528EEC9EFEB8AE8E318 /* ActivationIndicatorController.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ActivationIndicatorController.swift; sourceTree = ""; }; + D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ConstrainedBeamSearchTests.swift; sourceTree = ""; }; D9C1C921A1CDA2ADFC39EA01 /* AppsPaneView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = AppsPaneView.swift; sourceTree = ""; }; DB0CE9AB1286367BA2E82392 /* SettingsContainerView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SettingsContainerView.swift; sourceTree = ""; }; DDE858CB1E687E3CEB8FDD5B /* SuggestionRequestFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = SuggestionRequestFactory.swift; sourceTree = ""; }; @@ -741,6 +745,7 @@ 90B0D133AB77A2503FB08827 /* ClipboardRelevanceFilterTests.swift */, D504BEB224E0C176F5FCFF6E /* CompletionRenderModePolicyTests.swift */, 06FF2B0A3094A952A8EBA9B5 /* ConfidenceSuppressionPolicyTests.swift */, + D891397DAFCEF9DDD612A771 /* ConstrainedBeamSearchTests.swift */, 6B0200AA5C4A644A3B52A3EC /* ConstrainedSamplerTests.swift */, AF1E065C7FFB697FCEB2FA5C /* CotabbyTestFixtures.swift */, AD752451330486FE270018B0 /* CustomRulesTests.swift */, @@ -904,6 +909,7 @@ D3A2AC525DC664DB540D4F19 /* ClipboardRelevanceFilter.swift */, 53CF416511099C6818110F01 /* CompletionRenderModePolicy.swift */, 1BD71ECC2AE4821B643E0935 /* ConfidenceSuppressionPolicy.swift */, + A83264218423B61E59109EB3 /* ConstrainedBeamSearch.swift */, 47CF010A66DA31989524FCD0 /* ConstrainedSampler.swift */, C7B2D34A6F3AC9DFD61350F7 /* CotabbyDebugOptions.swift */, 29ED42C4BDD0C521101AF95E /* DeviceInfo.swift */, @@ -1099,6 +1105,7 @@ 7C94725B4837DEC9ECF1BC54 /* CompletionRenderMode.swift in Sources */, 3985F0F2B3178DBB945B1064 /* CompletionRenderModePolicy.swift in Sources */, 429CE592897D8A952F2916C3 /* ConfidenceSuppressionPolicy.swift in Sources */, + C638466A2F373BABA6F60A6D /* ConstrainedBeamSearch.swift in Sources */, 81FC391E375B7EEFF965FF1B /* ConstrainedSampler.swift in Sources */, 8B2DFC860803C0A7C4D34A36 /* ContextBuffer.swift in Sources */, AA2E09FF7E430D66ECA8ECD5 /* CotabbyApp.swift in Sources */, @@ -1262,6 +1269,7 @@ BFCA7FAFDAEBF586AB615567 /* ClipboardRelevanceFilterTests.swift in Sources */, 25F91CEF38400FD1ADB6B1AF /* CompletionRenderModePolicyTests.swift in Sources */, 91D8189EFCD1BA992EA6F038 /* ConfidenceSuppressionPolicyTests.swift in Sources */, + 29EC35D67D9B4C3C50222619 /* ConstrainedBeamSearchTests.swift in Sources */, 1C00CE3D553B58723EAE5F92 /* ConstrainedSamplerTests.swift in Sources */, 5E10EFC426217CB7218A5847 /* CotabbyTestFixtures.swift in Sources */, 91D1F16B8C5DA281D4B7F699 /* CustomRulesTests.swift in Sources */, diff --git a/Cotabby/Models/LlamaRuntimeModels.swift b/Cotabby/Models/LlamaRuntimeModels.swift index 7bc95945..af95ad95 100644 --- a/Cotabby/Models/LlamaRuntimeModels.swift +++ b/Cotabby/Models/LlamaRuntimeModels.swift @@ -199,6 +199,12 @@ struct LlamaGenerationOptions: Equatable, Sendable { /// validated on device. Changing it does not affect KV reuse, so it is intentionally excluded /// from `SamplingFingerprint`. var useConstrainedDecoder: Bool = false + + /// Beam width for the constrained decoder. 1 keeps the single-path greedy decode; values > 1 run a + /// multi-branch beam search that explores several short continuations and keeps the highest-scoring + /// one. Only consulted when `useConstrainedDecoder` is true. Like `useConstrainedDecoder`, it does + /// not affect KV reuse, so it is excluded from `SamplingFingerprint`. + var beamWidth: Int = 1 } /// The concrete runtime assets selected during bootstrap after checking available model files. diff --git a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift index 5b42873f..9a1b683d 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift @@ -198,9 +198,16 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { // The KV-trim defer above runs after whichever decoder returns. Both decoders share the // prepared sequence and the same confidence-suppression contract; they differ only in how // they pick each token (engine sampler vs. deterministic constrained selection). - return options.useConstrainedDecoder - ? try runConstrainedDecode(sequenceID: sequenceID, options: options) - : runEngineSampledDecode(sequenceID: sequenceID, options: options) + guard options.useConstrainedDecoder else { + return runEngineSampledDecode(sequenceID: sequenceID, options: options) + } + return options.beamWidth > 1 + ? try runConstrainedBeamDecode( + sequenceID: sequenceID, + promptTokenCount: promptTokens.count, + options: options + ) + : try runConstrainedDecode(sequenceID: sequenceID, options: options) } // MARK: - Decoders @@ -406,6 +413,128 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { return generatedText } + /// Multi-branch (beam) variant of the constrained decoder. Explores several short continuations + /// over the shared sequence — `EngineBeamStepper` syncs the KV cache to each branch's token path — + /// and returns the highest-scoring one. Reuses the same token profile, no-repeat-ngram guard, + /// sentence-boundary stop, and confidence suppression as the greedy path. The caller's KV-trim + /// defer restores the prompt-only state afterward. + private func runConstrainedBeamDecode( + sequenceID: Int32, + promptTokenCount: Int, + options: LlamaGenerationOptions + ) throws -> String { + let profile = try autocompleteTokenProfile() + let vocabSize = profile.vocabSize + guard vocabSize > 0 else { + throw LlamaRuntimeError.generationFailed("Vocabulary unavailable for constrained decoding.") + } + let topK = options.topK > 0 ? options.topK : vocabSize + var currentPath: [Int] = [] + var logitsBuffer = [Float](repeating: 0, count: vocabSize) + let candidates = ConstrainedBeamSearch.search( + nextLogits: { generatedTokens in + self.beamLogits( + forGeneratedTokens: generatedTokens, + sequenceID: sequenceID, + promptTokenCount: promptTokenCount, + currentPath: ¤tPath, + logitsBuffer: &logitsBuffer + ) + }, + profile: profile, + configuration: BeamSearchConfiguration( + beamWidth: options.beamWidth, + maxTokens: options.maxPredictionTokens, + topK: topK, + noRepeatNgramSize: Self.noRepeatNgramSize + ), + isSingleLine: options.singleLine + ) + let best = candidates.first + CotabbyLogger.runtime.debug( + "Decode end", + metadata: [ + "kind": .string("generate_beam"), + "beam_width": .stringConvertible(options.beamWidth), + "candidates": .stringConvertible(candidates.count), + "tokens_generated": .stringConvertible(best?.tokenIDs.count ?? 0) + ] + ) + guard let best else { + return "" + } + if Self.shouldSuppress( + sumLogprob: best.cumulativeLogprob, + tokensGenerated: best.tokenIDs.count, + options: options + ) { + return "" + } + return best.text + } + + /// Beam-search logits provider: syncs the shared sequence's KV to `generatedTokens`, then reads + /// the next-token logits. `currentPath` / `logitsBuffer` are owned by one beam run (the caller), + /// so this stays a plain method on the runtime where `engine` (a noncopyable C++ value) is mutable. + private func beamLogits( + forGeneratedTokens generatedTokens: [Int], + sequenceID: Int32, + promptTokenCount: Int, + currentPath: inout [Int], + logitsBuffer: inout [Float] + ) -> [Float]? { + guard syncBeamSequence( + to: generatedTokens, + sequenceID: sequenceID, + promptTokenCount: promptTokenCount, + currentPath: ¤tPath + ) else { + return nil + } + let vocabSize = logitsBuffer.count + let written = logitsBuffer.withUnsafeMutableBufferPointer { buffer in + Int(engine.getNextTokenLogits(sequenceID, buffer.baseAddress, Int32(buffer.count))) + } + guard written == vocabSize else { + return nil + } + return logitsBuffer + } + + /// Brings the sequence KV to exactly `target` tokens beyond the prompt: trim back to the longest + /// shared prefix with the current path, then accept the remaining target tokens. `currentPath` is + /// updated as tokens are accepted so it always reflects the real KV length, even on a mid-accept + /// failure (the caller treats a false return as "this branch cannot be extended"). + private func syncBeamSequence( + to target: [Int], + sequenceID: Int32, + promptTokenCount: Int, + currentPath: inout [Int] + ) -> Bool { + let shared = Self.commonPrefixLength(currentPath, target) + if currentPath.count > shared, !engine.trimKV(sequenceID, Int32(promptTokenCount + shared)) { + currentPath = [] + return false + } + currentPath = Array(target[.. Int { + var count = 0 + let limit = min(lhs.count, rhs.count) + while count < limit, lhs[count] == rhs[count] { + count += 1 + } + return count + } + /// Shared low-confidence gate for both decoders: drop completions the model itself was unsure /// about. Disabled by default (confidenceFloor == -infinity). The KV-trim defer in `generate` /// still runs because the caller returns "" rather than throwing. diff --git a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift index 0f084dba..4ada7c28 100644 --- a/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift +++ b/Cotabby/Services/Runtime/LlamaSuggestionEngine.swift @@ -22,6 +22,15 @@ final class LlamaSuggestionEngine { UserDefaults.standard.bool(forKey: constrainedDecoderDefaultsKey) } + /// UserDefaults key (no UI) for the constrained decoder's beam width. Default 1 keeps the existing + /// single-path greedy decode; a value > 1 runs a multi-branch beam search. Paired with the + /// constrained-decoder flag as a hidden developer/dogfood knob until validated on device. + private static let constrainedBeamWidthDefaultsKey = "cotabbyConstrainedBeamWidth" + private static var constrainedBeamWidth: Int { + let stored = UserDefaults.standard.integer(forKey: constrainedBeamWidthDefaultsKey) + return stored > 0 ? stored : 1 + } + init(runtimeManager: LlamaRuntimeGenerating) { self.runtimeManager = runtimeManager } @@ -60,7 +69,8 @@ final class LlamaSuggestionEngine { precedingText: request.context.precedingText, trailingText: request.context.trailingText ), - useConstrainedDecoder: Self.isConstrainedDecoderEnabled + useConstrainedDecoder: Self.isConstrainedDecoderEnabled, + beamWidth: Self.constrainedBeamWidth ) ) try Task.checkCancellation() diff --git a/Cotabby/Support/ConstrainedBeamSearch.swift b/Cotabby/Support/ConstrainedBeamSearch.swift new file mode 100644 index 00000000..a1669c57 --- /dev/null +++ b/Cotabby/Support/ConstrainedBeamSearch.swift @@ -0,0 +1,181 @@ +import Foundation + +/// File overview: +/// A deterministic multi-branch (beam) search over a model's next-token logits, used by the +/// constrained decoder to explore several short continuations at once and keep the highest-scoring +/// one instead of committing to a single greedy token at each step. +/// +/// Why this file exists: +/// Greedy argmax can paint a completion into a corner: a locally-best token that leads nowhere good. +/// A small beam keeps the best few branches per step and scores whole continuations by cumulative +/// log-probability, which yields steadier short completions. The algorithm is written against the +/// `BeamDecodeStepping` protocol so the search can be unit-tested against an in-memory fake that +/// returns fixed logits per token path; the live llama adapter is a thin KV-sync wrapper supplied by +/// the runtime. Everything here is pure and deterministic given the stepping inputs. + +/// The single operation a beam search needs from the model: the next-token logits for a branch's +/// token path (the tokens generated so far, beyond the prompt). The provider syncs its own KV state +/// to `generatedTokens` before reading and returns nil when logits are unavailable. Modeled as a +/// closure rather than a protocol so the live adapter can be an instance method on the runtime — the +/// inference engine is a noncopyable C++ type that cannot be stored in a separate object — while +/// tests pass a scripted closure. +typealias BeamLogitsProvider = (_ generatedTokens: [Int]) -> [Float]? + +/// Tuning for one beam search. `beamWidth` is the number of live branches kept per step; `topK` bounds +/// the per-branch expansion pool; `noRepeatNgramSize` forbids re-emitting an n-gram already in a +/// branch (the same anti-loop guard the greedy decoder uses). +struct BeamSearchConfiguration: Equatable { + let beamWidth: Int + let maxTokens: Int + let topK: Int + let noRepeatNgramSize: Int + + init(beamWidth: Int, maxTokens: Int, topK: Int, noRepeatNgramSize: Int = 3) { + self.beamWidth = max(1, beamWidth) + self.maxTokens = max(0, maxTokens) + self.topK = topK + self.noRepeatNgramSize = noRepeatNgramSize + } +} + +/// One explored continuation: the committed token ids (beyond the prompt), their accumulated raw +/// UTF-8 bytes, and the cumulative log-probability of the path. +struct BeamCandidate: Equatable { + let tokenIDs: [Int] + let bytes: [UInt8] + let cumulativeLogprob: Double + + /// Mean per-token log-probability; ranks completed branches so a short confident continuation is + /// not unfairly beaten by a longer, lower-average one. An empty branch ranks last. + var meanLogprob: Double { + tokenIDs.isEmpty ? -.infinity : cumulativeLogprob / Double(tokenIDs.count) + } + + /// The decoded text of the branch. Lossy on a partial trailing scalar (a final token may carry + /// only part of a multi-byte character), which renders as U+FFFD rather than dropping the branch. + var text: String { + // swiftlint:disable:next optional_data_string_conversion + String(decoding: bytes, as: UTF8.self) + } +} + +enum ConstrainedBeamSearch { + /// Runs a best-first beam search and returns completed branches, best mean-logprob first. A branch + /// completes on an end-of-generation token, a single-line newline (not emitted), a sentence + /// boundary, the token budget, or when no admissible token remains. Pure given `stepping`. + static func search( + nextLogits: @escaping BeamLogitsProvider, + profile: TokenProfile, + configuration: BeamSearchConfiguration, + isSingleLine: Bool + ) -> [BeamCandidate] { + Engine( + nextLogits: nextLogits, + profile: profile, + configuration: configuration, + isSingleLine: isSingleLine + ).run() + } +} + +/// The mutable-free search context, holding the inputs so the per-step helpers stay small. A struct +/// rather than passing the same four values through every call. +private struct Engine { + let nextLogits: BeamLogitsProvider + let profile: TokenProfile + let configuration: BeamSearchConfiguration + let isSingleLine: Bool + + func run() -> [BeamCandidate] { + var frontier: [BeamCandidate] = [BeamCandidate(tokenIDs: [], bytes: [], cumulativeLogprob: 0)] + var completed: [BeamCandidate] = [] + + for _ in 0 ..< configuration.maxTokens { + guard !frontier.isEmpty else { break } + var nextFrontier: [BeamCandidate] = [] + for branch in frontier { + guard let logits = nextLogits(branch.tokenIDs) else { + completed.append(branch) + continue + } + expand(branch: branch, logits: logits, live: &nextFrontier, completed: &completed) + } + frontier = Self.prune(nextFrontier, to: configuration.beamWidth) + } + // Budget exhausted: surviving branches are valid completions too. + completed.append(contentsOf: frontier) + return completed + .filter { !$0.tokenIDs.isEmpty } + .sorted { $0.meanLogprob > $1.meanLogprob } + } + + /// Expands one branch across its admissible tokens, routing branch-ending tokens (end-of- + /// generation, single-line newline, sentence boundary) into `completed` and the rest into `live`. + private func expand( + branch: BeamCandidate, + logits: [Float], + live: inout [BeamCandidate], + completed: inout [BeamCandidate] + ) { + let blocked = RepetitionGuard.blockedTokens( + history: branch.tokenIDs, + ngramSize: configuration.noRepeatNgramSize + ) + let candidates = ConstrainedSampler.rankedAdmissibleTokens( + logits: logits, + profile: profile, + admissibleTokenIDs: nil, + topK: configuration.topK, + blockedTokenIDs: blocked + ) + for tokenID in candidates { + if profile.isEndOfGeneration(tokenID) { + completed.append(branch) + continue + } + if isSingleLine, profile.isNewline(tokenID) { + completed.append(branch) + continue + } + let tokenBytes = profile.bytes(for: tokenID) + let child = extend(branch, by: tokenID, tokenBytes: tokenBytes, logits: logits) + if Self.completesSentence(child.bytes, lastTokenBytes: tokenBytes) { + completed.append(child) + } else { + live.append(child) + } + } + } + + private func extend( + _ branch: BeamCandidate, + by tokenID: Int, + tokenBytes: [UInt8], + logits: [Float] + ) -> BeamCandidate { + var bytes = branch.bytes + bytes.append(contentsOf: tokenBytes) + let logprob = ConstrainedSampler.logProb(ofTokenAt: tokenID, in: logits) ?? 0 + return BeamCandidate( + tokenIDs: branch.tokenIDs + [tokenID], + bytes: bytes, + cumulativeLogprob: branch.cumulativeLogprob + logprob + ) + } + + /// Keeps the best `width` live branches by cumulative log-probability. + private static func prune(_ branches: [BeamCandidate], to width: Int) -> [BeamCandidate] { + guard branches.count > width else { + return branches + } + return Array(branches.sorted { $0.cumulativeLogprob > $1.cumulativeLogprob }.prefix(width)) + } + + private static func completesSentence(_ bytes: [UInt8], lastTokenBytes: [UInt8]) -> Bool { + guard lastTokenBytes.contains(where: { $0 == 0x2E || $0 == 0x21 || $0 == 0x3F }) else { + return false + } + // swiftlint:disable:next optional_data_string_conversion + return SentenceBoundaryClassifier.endsSentence(String(decoding: bytes, as: UTF8.self)) + } +} diff --git a/Cotabby/Support/ConstrainedSampler.swift b/Cotabby/Support/ConstrainedSampler.swift index 1e6949ae..5a830131 100644 --- a/Cotabby/Support/ConstrainedSampler.swift +++ b/Cotabby/Support/ConstrainedSampler.swift @@ -72,6 +72,38 @@ enum ConstrainedSampler { return best } + /// The admissible token ids for a step, ranked highest-logit first. Survivors are the same set + /// `selectToken` would consider — in-range, not `profile.isExcluded`, not in `blockedTokenIDs`, + /// and, when `admissibleTokenIDs` is non-nil, members of that set — and at most `topK` are + /// returned. This is the multi-candidate form of `selectToken`: the beam search expands a branch + /// across these instead of committing to the single best. Ties break by lower id for determinism. + static func rankedAdmissibleTokens( + logits: [Float], + profile: TokenProfile, + admissibleTokenIDs: Set?, + topK: Int, + blockedTokenIDs: Set = [] + ) -> [Int] { + guard topK > 0, !logits.isEmpty else { + return [] + } + if let admissible = admissibleTokenIDs, admissible.isEmpty { + return [] + } + let survivors = (0 ..< logits.count).filter { id in + !profile.isExcluded(id) + && !blockedTokenIDs.contains(id) + && (admissibleTokenIDs?.contains(id) ?? true) + } + let ranked = survivors.sorted { lhs, rhs in + if logits[lhs] != logits[rhs] { + return logits[lhs] > logits[rhs] + } + return lhs < rhs + } + return Array(ranked.prefix(topK)) + } + /// Average per-step log-probability of a sequence of chosen tokens, a confidence summary suitable /// for the existing low-confidence suppression policy. /// diff --git a/CotabbyTests/ConstrainedBeamSearchTests.swift b/CotabbyTests/ConstrainedBeamSearchTests.swift new file mode 100644 index 00000000..394443ef --- /dev/null +++ b/CotabbyTests/ConstrainedBeamSearchTests.swift @@ -0,0 +1,164 @@ +import XCTest +@testable import Cotabby + +/// Tests for the constrained beam search, exercised against a scripted `BeamLogitsProvider` that maps +/// a generated-token path to a logits row. This validates the whole search algorithm deterministically +/// without loading a model; the live llama adapter is a thin KV-sync method on the runtime. +final class ConstrainedBeamSearchTests: XCTestCase { + + // MARK: - Fixtures + + /// Records the token paths the search queried, so a test can assert it stopped where expected. + private final class PathRecorder { + private(set) var paths: [[Int]] = [] + func record(_ path: [Int]) { paths.append(path) } + } + + /// A scripted logits provider: returns the row mapped for a path, or a uniform low row otherwise. + private func provider( + vocabSize: Int, + rows: [[Int]: [Float]], + recorder: PathRecorder? = nil + ) -> BeamLogitsProvider { + { path in + recorder?.record(path) + return rows[path] ?? [Float](repeating: -20, count: vocabSize) + } + } + + private func makeProfile(byteStrings: [String], eog: Set = []) -> TokenProfile { + let bytes = byteStrings.map { Array($0.utf8) } + return TokenProfile.build( + vocabSize: bytes.count, + bytesFor: { bytes[$0] }, + isControl: { bytes[$0].isEmpty }, + isEndOfGeneration: { eog.contains($0) } + ) + } + + private func row(_ values: [Int: Float], vocabSize: Int) -> [Float] { + var logits = [Float](repeating: -20, count: vocabSize) + for (id, value) in values { + logits[id] = value + } + return logits + } + + // MARK: - rankedAdmissibleTokens + + func test_rankedAdmissibleTokens_ordersByLogitAndDropsExcludedAndBlocked() { + // token 2 is control (empty bytes) -> excluded; token 0 is blocked by the caller. + let profile = makeProfile(byteStrings: ["a", "b", "", "d"]) + let logits: [Float] = [9, 8, 7, 6] + + let ranked = ConstrainedSampler.rankedAdmissibleTokens( + logits: logits, profile: profile, admissibleTokenIDs: nil, topK: 10, blockedTokenIDs: [0]) + + XCTAssertEqual(ranked, [1, 3], "drops blocked 0 and control 2; orders the rest by logit") + } + + func test_rankedAdmissibleTokens_capsAtTopK() { + let profile = makeProfile(byteStrings: ["a", "b", "c", "d"]) + let ranked = ConstrainedSampler.rankedAdmissibleTokens( + logits: [1, 4, 3, 2], profile: profile, admissibleTokenIDs: nil, topK: 2) + XCTAssertEqual(ranked, [1, 2]) + } + + // MARK: - search + + func test_search_widthOne_takesTheHighestLogitToken() { + let profile = makeProfile(byteStrings: ["a", "b"]) + let result = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 2, rows: [[]: row([0: 2, 1: 1], vocabSize: 2)]), + profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 1, maxTokens: 1, topK: 5), + isSingleLine: false) + + XCTAssertEqual(result.first?.text, "a") + XCTAssertEqual(result.first?.tokenIDs, [0]) + } + + func test_search_stopsOnEndOfGenerationToken() { + // Token 2 is EOG with non-empty bytes (so it is not masked as a control token). When it is the + // chosen next token, the branch completes without emitting it. (Real EOG tokens render empty + // and are excluded like any control token; this drives the EOG branch directly.) + let profile = makeProfile(byteStrings: ["a", "b", "Z"], eog: [2]) + let result = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 3, rows: [ + []: row([0: 5], vocabSize: 3), + [0]: row([2: 9], vocabSize: 3) + ]), + profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 1, maxTokens: 4, topK: 2), + isSingleLine: false) + + XCTAssertEqual(result.first?.text, "a") + XCTAssertEqual(result.first?.tokenIDs, [0]) + } + + func test_search_widerBeamExploresASecondBestFirstTokenThatGreedyMisses() { + // 0="A" (best first), 1="B" (second). Each leads to a continuation that ends a sentence, so the + // branches complete cleanly. Greedy only follows "A"; a width-2 beam also reaches "By.". + let profile = makeProfile(byteStrings: ["A", "B", "x.", "y."]) + let rows: [[Int]: [Float]] = [ + []: row([0: 5, 1: 4], vocabSize: 4), + [0]: row([2: 9], vocabSize: 4), + [1]: row([3: 9], vocabSize: 4) + ] + let greedy = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 4, rows: rows), profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 1, maxTokens: 4, topK: 2), + isSingleLine: false) + let beam = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 4, rows: rows), profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 2, maxTokens: 4, topK: 2), + isSingleLine: false) + + XCTAssertFalse(greedy.map(\.text).contains("By."), "greedy only explores the best first token") + XCTAssertTrue(beam.map(\.text).contains("By."), "a width-2 beam also explores the second-best first token") + } + + func test_search_doesNotEmitNewlineInSingleLineField() { + // token 0 is a newline, token 1 is "a". In a single-line field the newline ends a branch + // without being emitted. + let profile = makeProfile(byteStrings: ["\n", "a"]) + let result = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 2, rows: [[]: row([0: 5, 1: 4], vocabSize: 2)]), + profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 2, maxTokens: 1, topK: 5), + isSingleLine: true) + + XCTAssertEqual(result.first?.text, "a") + XCTAssertFalse(result.contains { $0.text.contains("\n") }) + } + + func test_search_stopsAtSentenceBoundaryAndDoesNotGenerateFurther() { + // 0="Done", 1=".", 2="More". The completion should stop after the period and never step past. + let profile = makeProfile(byteStrings: ["Done", ".", "More"]) + let recorder = PathRecorder() + let result = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 3, rows: [ + []: row([0: 5], vocabSize: 3), + [0]: row([1: 5], vocabSize: 3), + [0, 1]: row([2: 5], vocabSize: 3) + ], recorder: recorder), + profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 1, maxTokens: 6, topK: 5), + isSingleLine: false) + + XCTAssertEqual(result.first?.text, "Done.") + XCTAssertFalse(recorder.paths.contains([0, 1]), "search must stop at the sentence and not step past it") + } + + func test_search_respectsMaxTokenBudget() { + // No EOG / sentence end: every token keeps generating, so the budget bounds the length. + let profile = makeProfile(byteStrings: ["a", "b"]) + let result = ConstrainedBeamSearch.search( + nextLogits: provider(vocabSize: 2, rows: [:]), // unmapped -> uniform low logits + profile: profile, + configuration: BeamSearchConfiguration(beamWidth: 1, maxTokens: 2, topK: 2), + isSingleLine: false) + + XCTAssertEqual(result.first?.tokenIDs.count, 2) + } +}