Stop the constrained decoder at sentence boundaries

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,30 +355,33 @@ 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
 
             if engine.acceptToken(sequenceID, Int32(tokenID)) != .ok {
                 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

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}"
+    }
+}

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(""))
+    }
 }