From 7fb88ede48eef7c853c058cd818cea8682cd4d78 Mon Sep 17 00:00:00 2001 From: Mason363 Date: Tue, 8 Sep 2026 22:55:52 -0400 Subject: [PATCH] Add a required-prefix constraint for word completion setRequiredPrefix(sequence, bytes) makes the next generation begin with the given bytes: while any remain, every token whose text is inconsistent with them is masked before sampling, and each sampled token consumes what it covers. A caller completing a half-typed word prompts up to the word boundary, so the tokenizer sees whole words rather than a word cut at an arbitrary byte, and requires the completion to start with the boundary whitespace plus the typed letters. Token texts are cached at model load so the mask costs one pass over the vocabulary per constrained token; a row with no consistent token drops the constraint instead of sampling from an all-masked distribution. --- .../CotabbyInferenceEngine.cpp | 72 +++++++++++++++++++ .../include/CotabbyInferenceEngine.h | 10 +++ .../LlamaMiddlewareTests.swift | 36 ++++++++++ 3 files changed, 118 insertions(+) diff --git a/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp b/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp index 5fc867f..2e287a6 100644 --- a/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp +++ b/Sources/CotabbyInferenceEngine/CotabbyInferenceEngine.cpp @@ -112,6 +112,10 @@ struct SequenceState { // Set by setForceWordContinuation; consumed (and cleared) when the next seed token is sampled. bool force_word_continuation = false; + // Bytes the generation must still produce before it runs free (see setRequiredPrefix). Each + // sampled token's text is consumed from the front; empty means unconstrained. + std::string required_prefix; + // Whether computeLogprob runs for this sequence's tokens. Defaults to true (the historical // behavior) so existing callers keep getting real log-probabilities; callers whose confidence // gate is disabled opt out via setComputeLogprob to skip two O(vocab) passes per token. @@ -154,6 +158,9 @@ struct CotabbyInferenceEngine::Impl { std::vector nonprintable_bias; std::vector linebreak_bias; std::vector starts_new_word; + // Every token's plain-rendered text, so a required-prefix mask can compare bytes without + // calling the tokenizer per token on the hot path. Empty for tokens that render to nothing. + std::vector token_pieces; // One product sequence with a monotonically changing external identity. The mutex protects // create/destroy and lookup; callers still must not destroy the sequence while another method @@ -250,6 +257,7 @@ struct CotabbyInferenceEngine::Impl { const int32_t n = llama_vocab_n_tokens(vocab); starts_new_word.assign(static_cast(n), false); + token_pieces.assign(static_cast(n), std::string()); // BOS belongs at sequence start only; some vocabularies ship it without the control // attribute, which would otherwise let it be sampled mid-text. @@ -286,6 +294,7 @@ struct CotabbyInferenceEngine::Impl { if (written <= 0) { continue; } + token_pieces[static_cast(t)].assign(piece, static_cast(written)); const char first = piece[0]; if (first == ' ' || first == '\t' || first == '\n' || first == '\r') { starts_new_word[static_cast(t)] = true; @@ -315,6 +324,48 @@ struct CotabbyInferenceEngine::Impl { } } + // Masks every token whose text cannot begin the still-required bytes `remaining`: a token is + // consistent when its text is a prefix of `remaining` (it consumes part of it) or `remaining` is + // a prefix of its text (it consumes all of it and continues). Tokens that render to nothing, + // EOG among them, are never consistent. Returns false, masking nothing, in the pathological + // case where no token is consistent, so the sampler is never handed an all-masked row. + bool maskInconsistentWithPrefix(int logits_row, const std::string& remaining) { + if (!shared_ctx || !vocab || remaining.empty()) return false; + float* logits = llama_get_logits_ith(shared_ctx, logits_row); + if (!logits) return false; + const size_t n = token_pieces.size(); + std::vector consistent(n, false); + size_t consistent_count = 0; + for (size_t t = 0; t < n; ++t) { + const std::string& piece = token_pieces[t]; + if (piece.empty()) continue; + const size_t overlap = std::min(piece.size(), remaining.size()); + if (remaining.compare(0, overlap, piece, 0, overlap) == 0) { + consistent[t] = true; + ++consistent_count; + } + } + if (consistent_count == 0) return false; + for (size_t t = 0; t < n; ++t) { + if (!consistent[t]) { + logits[t] = -INFINITY; + } + } + return true; + } + + // Consumes a sampled token's text from the sequence's required prefix. + void consumeRequiredPrefix(SequenceState* seq, llama_token token) const { + if (!seq || seq->required_prefix.empty()) return; + const size_t index = static_cast(token); + const size_t consumed = index < token_pieces.size() ? token_pieces[index].size() : 0; + if (consumed >= seq->required_prefix.size()) { + seq->required_prefix.clear(); + } else { + seq->required_prefix.erase(0, consumed); + } + } + // Log-probability of `token` under the raw model distribution at `logits_row`, used as a // confidence signal. Two O(vocab) passes; only invoked on the autocomplete path. float computeLogprob(int logits_row, llama_token token) const { @@ -629,11 +680,17 @@ EngineStatus CotabbyInferenceEngine::decodePrompt(int32_t sequence_id, impl_->maskNewWordStarts(-1); seq->force_word_continuation = false; } + // Required-prefix constraint (see setRequiredPrefix): the seed must begin the required bytes. + // A row with no consistent token drops the constraint rather than sampling from nothing. + if (!seq->required_prefix.empty() && !impl_->maskInconsistentWithPrefix(-1, seq->required_prefix)) { + seq->required_prefix.clear(); + } // Seed sample: take one token from the prompt's final logits row. The seed will be returned by // the next sampleNext call as-is and feedback-decoded by the call after that. llama_token seed = llama_sampler_sample(seq->sampler, impl_->shared_ctx, -1); llama_sampler_accept(seq->sampler, seed); + impl_->consumeRequiredPrefix(seq, seed); seq->seed_token = seed; seq->seed_logprob = seq->compute_logprob ? impl_->computeLogprob(-1, seed) : 0.0f; seq->seed_argmax_is_eog = impl_->argmaxIsEOG(-1); @@ -737,7 +794,11 @@ SampleResult CotabbyInferenceEngine::sampleNext(int32_t sequence_id) { } else if (seq->cancelled.load(std::memory_order_acquire)) { result.was_cancelled = true; } else { + if (!seq->required_prefix.empty() && !impl_->maskInconsistentWithPrefix(0, seq->required_prefix)) { + seq->required_prefix.clear(); + } const llama_token next = llama_sampler_sample(seq->sampler, impl_->shared_ctx, 0); + impl_->consumeRequiredPrefix(seq, next); result.argmax_is_eog = impl_->argmaxIsEOG(0); result.token = next; @@ -822,6 +883,17 @@ void CotabbyInferenceEngine::setForceWordContinuation(int32_t sequence_id, bool } } +void CotabbyInferenceEngine::setRequiredPrefix(int32_t sequence_id, const char* utf8, int length) { + if (!impl_) return; + SequenceState* seq = impl_->findSequence(sequence_id); + if (!seq) return; + if (!utf8 || length <= 0) { + seq->required_prefix.clear(); + return; + } + seq->required_prefix.assign(utf8, static_cast(length)); +} + void CotabbyInferenceEngine::setComputeLogprob(int32_t sequence_id, bool enabled) { if (!impl_) return; SequenceState* seq = impl_->findSequence(sequence_id); diff --git a/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h b/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h index 383bcaf..14babe8 100644 --- a/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h +++ b/Sources/CotabbyInferenceEngine/include/CotabbyInferenceEngine.h @@ -80,6 +80,16 @@ class CotabbyInferenceEngine { // the constraint clears. Set this before `decodePrompt`, which samples the first (seed) token. void setForceWordContinuation(int32_t sequence_id, bool enabled); + // Constrains the generation on `sequence_id` to begin with the UTF-8 bytes `utf8[0..length)`. + // While any of them are unconsumed, only tokens whose text is a prefix of the remainder, or + // that the remainder is a prefix of, can be sampled; the remainder shrinks by each sampled + // token's text and the constraint clears once it is empty. This is how a caller completes a + // word the user has half typed: prompt up to the word boundary (so the tokenizer sees whole + // words, not a word cut at an arbitrary byte) and require the completion to start with the + // boundary whitespace plus the typed letters. Set before `decodePrompt`, which samples the + // first token; an empty prefix clears any pending constraint. + void setRequiredPrefix(int32_t sequence_id, const char* utf8, int length); + // Controls whether `SampleResult.logprob` is computed for this sequence. Defaults to true // (the historical behavior). The log-probability costs two O(vocab-size) passes per generated // token, so callers whose confidence gating is disabled should pass false to skip it; results diff --git a/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift b/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift index b7741ba..da73cb4 100644 --- a/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift +++ b/Tests/CotabbyInferenceTests/LlamaMiddlewareTests.swift @@ -149,6 +149,42 @@ final class LlamaMiddlewareTests: XCTestCase { engine.destroySequence(sequence) } + func testRequiredPrefixConstrainsTheStartOfTheCompletion() throws { + let modelPath = try Self.modelPath() + var engine = CotabbyInferenceEngine() + XCTAssertEqual(engine.loadModel(modelPath, -1, 1024, 256), EngineStatus.ok) + defer { engine.unloadModel() } + + // The user has typed "Thanks for sending over the draft yest": the prompt stops at the word + // boundary and the completion must begin with the boundary space plus the typed letters, so + // the model completes the word the user started instead of a word cut at a token boundary. + let sequence = engine.createSequence(Self.samplingConfig(temperature: 0)) + let prompt = "Thanks for sending over the draft" + var tokens = Array(engine.tokenize(prompt, Int32(prompt.utf8.count))) + let required = " yest" + required.withCString { engine.setRequiredPrefix(sequence, $0, Int32(required.utf8.count)) } + XCTAssertEqual( + engine.decodePrompt(sequence, &tokens, Int32(tokens.count), 0), + EngineStatus.ok + ) + + var text = "" + for _ in 0 ..< 12 { + let result = engine.sampleNext(sequence) + if result.is_eos || result.was_cancelled { break } + text += Self.string(from: result) + } + XCTAssertTrue(text.hasPrefix(required), "got \(text)") + XCTAssertTrue(text.hasPrefix(" yesterday"), "the model should finish the word: \(text)") + engine.destroySequence(sequence) + } + + func testRequiredPrefixIsClearedByAnEmptyPrefixAndIgnoredWithoutASequence() throws { + var engine = CotabbyInferenceEngine() + "abc".withCString { engine.setRequiredPrefix(999, $0, 3) } + engine.setRequiredPrefix(999, nil, 0) + } + func testSampleNextReportsFiniteLogprob() throws { let modelPath = try Self.modelPath() var engine = CotabbyInferenceEngine()