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()