diff --git a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift index b71e9679..5b42873f 100644 --- a/Cotabby/Services/Runtime/LlamaRuntimeCore.swift +++ b/Cotabby/Services/Runtime/LlamaRuntimeCore.swift @@ -210,6 +210,42 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { /// single-token runs after a few repeats, without blocking ordinary short repeats like "very very". private static let noRepeatNgramSize = 3 + /// ASCII sentence terminators (`.` `!` `?`), used as a cheap pre-filter: the constrained decoder + /// only decodes the accumulated bytes to test for a sentence boundary when a token carried one of + /// these, keeping the steady path on raw byte accumulation. + private static func isSentenceTerminatorByte(_ byte: UInt8) -> Bool { + byte == 0x2E || byte == 0x21 || byte == 0x3F + } + + /// The stop reason for a token that must end the loop *before* it is committed to the output: an + /// end-of-generation token, or a line break in a single-line field. Returns nil to commit the + /// token. Folding both checks into one helper keeps the decode loop under the complexity budget. + private static func preCommitStopReason( + tokenID: Int, + options: LlamaGenerationOptions, + profile: TokenProfile + ) -> String? { + if profile.isEndOfGeneration(tokenID) { + return "eos" + } + if options.singleLine, profile.isNewline(tokenID) { + return "single_line" + } + return nil + } + + /// Whether the accumulated completion bytes now end a sentence. Decodes only when the last token + /// carried an ASCII sentence terminator, so the steady decode path avoids per-token String work. + /// Kept as a helper so the decode loop stays under the cyclomatic-complexity budget. + private static func completesSentence(_ generatedBytes: [UInt8], lastTokenBytes: [UInt8]) -> Bool { + guard lastTokenBytes.contains(where: isSentenceTerminatorByte) else { + return false + } + // swiftlint:disable:next optional_data_string_conversion + let decoded = String(decoding: generatedBytes, as: UTF8.self) + return SentenceBoundaryClassifier.endsSentence(decoded) + } + /// The shipping decoder: delegates token selection to the engine's built-in sampler /// (`sampleNext`), which applies temperature / top-k / top-p / min-p and commits each token. private func runEngineSampledDecode(sequenceID: Int32, options: LlamaGenerationOptions) -> String { @@ -319,23 +355,20 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { break } - if profile.isEndOfGeneration(tokenID) { - stopReason = "eos" - break - } - // Single-line fields must never receive a line break; stop before emitting one so the - // partial completion so far is preserved (mirrors the sampler path's single_line mask). - if options.singleLine, profile.isNewline(tokenID) { - stopReason = "single_line" + // A token can stop the loop before it is committed: an end-of-generation token, or a line + // break in a single-line field (the partial completion so far is preserved). + if let preCommitStop = Self.preCommitStopReason(tokenID: tokenID, options: options, profile: profile) { + stopReason = preCommitStop break } // Accumulate raw bytes and decode once at the end: a single token may carry only part of // a multi-byte UTF-8 scalar, so per-token String decoding would corrupt CJK / emoji. + let tokenBytes = profile.bytes(for: tokenID) if let logProb = ConstrainedSampler.logProb(ofTokenAt: tokenID, in: logits) { sumLogprob += logProb } - generatedBytes.append(contentsOf: profile.bytes(for: tokenID)) + generatedBytes.append(contentsOf: tokenBytes) generatedTokenIDs.append(tokenID) tokensGenerated += 1 @@ -343,6 +376,12 @@ nonisolated final class LlamaRuntimeCore: @unchecked Sendable { stopReason = "accept_failed" break } + + // Stop cleanly at the end of a sentence rather than running into the next one. + if Self.completesSentence(generatedBytes, lastTokenBytes: tokenBytes) { + stopReason = "sentence_boundary" + break + } } // Lossy decode is deliberate: the accumulated bytes are valid UTF-8 except for a possible diff --git a/Cotabby/Support/SentenceBoundaryClassifier.swift b/Cotabby/Support/SentenceBoundaryClassifier.swift index 0bb82a92..7437f46d 100644 --- a/Cotabby/Support/SentenceBoundaryClassifier.swift +++ b/Cotabby/Support/SentenceBoundaryClassifier.swift @@ -15,6 +15,37 @@ enum SentenceBoundaryClassifier { "mr", "mrs", "ms", "dr", "st", "vs", "eg", "ie", "etc", "no", "fig", "approx", "inc", "ltd" ] + /// Whether `text` ends at a sentence boundary: after skipping any trailing whitespace and a run of + /// closing punctuation (quotes / brackets), the last visible character is `!`, `?`, or a *terminal* + /// period. The closing-punctuation skip lets `He said "stop."` and `(done!)` count as sentence ends + /// even though their final character is a quote or paren. Used to stop a constrained completion + /// cleanly at the end of one sentence instead of letting it run into the next. + static func endsSentence(_ text: String) -> Bool { + var index = text.endIndex + while index > text.startIndex { + let previous = text.index(before: index) + guard text[previous].isWhitespace else { break } + index = previous + } + while index > text.startIndex { + let previous = text.index(before: index) + guard text[previous].isSentenceClosingPunctuation else { break } + index = previous + } + guard index > text.startIndex else { + return false + } + let lastIndex = text.index(before: index) + switch text[lastIndex] { + case "!", "?": + return true + case ".": + return isTerminalPeriod(in: text, at: lastIndex) + default: + return false + } + } + /// Whether the period at `periodIndex` in `text` ends a sentence. The caller guarantees that /// `text[periodIndex]` is ".". static func isTerminalPeriod(in text: String, at periodIndex: String.Index) -> Bool { @@ -63,3 +94,13 @@ enum SentenceBoundaryClassifier { return String(letters.reversed()) } } + +private extension Character { + /// Closing punctuation that may follow a sentence terminator: straight and curly quotes, + /// parentheses, square brackets, and braces. `endsSentence` walks back past a run of these to find + /// the real terminator underneath, so `"done."` and `(stop!)` register as sentence ends. + var isSentenceClosingPunctuation: Bool { + self == "\"" || self == "'" || self == ")" || self == "]" || self == "}" + || self == "\u{201D}" || self == "\u{2019}" + } +} diff --git a/CotabbyTests/SentenceBoundaryClassifierTests.swift b/CotabbyTests/SentenceBoundaryClassifierTests.swift index 5c993823..f947a48e 100644 --- a/CotabbyTests/SentenceBoundaryClassifierTests.swift +++ b/CotabbyTests/SentenceBoundaryClassifierTests.swift @@ -41,4 +41,39 @@ final class SentenceBoundaryClassifierTests: XCTestCase { let text = "tabs and so on etc." XCTAssertFalse(SentenceBoundaryClassifier.isTerminalPeriod(in: text, at: lastPeriodIndex(in: text))) } + + // MARK: - endsSentence + + func test_endsSentence_trueForTerminalPeriod() { + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("Hello world.")) + } + + func test_endsSentence_falseWithoutTerminator() { + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("Hello world")) + } + + func test_endsSentence_trueForExclamationAndQuestion() { + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("Yes!")) + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("Really?")) + } + + func test_endsSentence_ignoresTrailingWhitespace() { + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("Done. ")) + } + + func test_endsSentence_walksPastClosingPunctuation() { + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("He said \"stop.\"")) + XCTAssertTrue(SentenceBoundaryClassifier.endsSentence("(done!)")) + } + + func test_endsSentence_falseForNonTerminalPeriods() { + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("It is version 1.")) + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("for example, e.g.")) + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("from the U.S.")) + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("Hello Mr.")) + } + + func test_endsSentence_falseForEmptyString() { + XCTAssertFalse(SentenceBoundaryClassifier.endsSentence("")) + } }