add check script

common/arg.cpp

@@ -2159,6 +2159,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex
             params.sampling.dry_penalty_last_n = value;
         }
     ).set_sampling());
+    add_opt(common_arg(
+        {"--synthid-keys"}, "KEYS",
+        "comma-separated list of secret SynthID watermarking keys, must match the keys used by the detector, compatible with HF transformers SynthIDTextWatermarkingConfig (default: disabled)",
+        [](common_params & params, const std::string & value) {
+            params.sampling.synthid_keys.clear();
+            for (const auto & key : string_split<std::string>(value, ',')) {
+                params.sampling.synthid_keys.push_back(std::stoll(key));
+            }
+            if (params.sampling.synthid_keys.empty()) {
+                throw std::runtime_error("error: --synthid-keys requires at least one key");
+            }
+        }
+    ).set_sampling().set_env("LLAMA_ARG_SYNTHID_KEYS"));
     add_opt(common_arg(
         {"--dry-sequence-breaker"}, "STRING",
         string_format("add sequence breaker for DRY sampling, clearing out default breakers (%s) in the process; use \"none\" to not use any sequence breakers\n",

common/common.h

@@ -258,6 +258,8 @@ struct common_params_sampling {
 
     std::vector<std::string> dry_sequence_breakers = {"\n", ":", "\"", "*"};     // default sequence breakers for DRY
 
+    std::vector<int64_t> synthid_keys; // SynthID watermarking keys (empty = disabled)
+
     std::vector<enum common_sampler_type> samplers = {
         COMMON_SAMPLER_TYPE_PENALTIES,
         COMMON_SAMPLER_TYPE_DRY,

common/sampling.cpp

@@ -12,6 +12,7 @@
 #include <climits>
 #include <cmath>
 #include <cstring>
+#include <random>
 #include <unordered_map>
 #include <vector>
 
@@ -184,6 +185,24 @@ std::string common_params_sampling::print() const {
     return std::string(result);
 }
 
+// same defaults as SynthIDTextWatermarkingConfig from HF transformers
+// the sampling table is generated the same way as torch.randint(0, 2, ...) with a seeded CPU generator
+static struct llama_sampler * common_sampler_init_synthid(const std::vector<int64_t> & keys) {
+    const int32_t ngram_len            = 5;
+    const int32_t context_history_size = 1024;
+    const size_t  sampling_table_size  = 65536;
+    const uint32_t sampling_table_seed = 0;
+
+    std::vector<uint8_t> sampling_table(sampling_table_size);
+
+    std::mt19937 rng(sampling_table_seed);
+    for (auto & v : sampling_table) {
+        v = rng() % 2;
+    }
+
+    return llama_sampler_init_synthid(keys.data(), keys.size(), sampling_table.data(), sampling_table.size(), ngram_len, context_history_size);
+}
+
 struct common_sampler * common_sampler_init(
         const struct llama_model * model,
         struct common_params_sampling & params) {
@@ -337,6 +356,10 @@ struct common_sampler * common_sampler_init(
         }
     }
 
+    if (params.mirostat != 0 && !params.synthid_keys.empty()) {
+        LOG_WRN("%s: SynthID watermarking is not supported with mirostat, disabling\n", __func__);
+    }
+
     if (params.mirostat == 0) {
 
         bool use_adaptive_p = false; // see below
@@ -391,6 +414,10 @@ struct common_sampler * common_sampler_init(
                     GGML_ASSERT(false && "unknown sampler type");
             }
         }
+        if (!params.synthid_keys.empty()) {
+            samplers.push_back(common_sampler_init_synthid(params.synthid_keys));
+        }
+
         if (use_adaptive_p) {
             // only if user explicitly included adaptive-p sampler
             samplers.push_back(llama_sampler_init_adaptive_p(params.adaptive_target, params.adaptive_decay, params.seed));

include/llama.h

@@ -1394,6 +1394,20 @@ extern "C" {
     /// @details Top n sigma sampling as described in academic paper "Top-nσ: Not All Logits Are You Need" https://arxiv.org/pdf/2411.07641
     LLAMA_API struct llama_sampler * llama_sampler_init_top_n_sigma(float   n);
 
+    /// @details SynthID text watermarking, compatible with the SynthIDTextWatermarkLogitsProcessor from HF transformers
+    /// place it after truncation and temperature samplers and before the dist sampler
+    /// @param keys secret watermarking keys, each key adds one tournament round (the HF default config uses 30 keys)
+    /// @param sampling_table table of 0/1 values, maps a hashed ngram to a g-value
+    /// @param ngram_len number of tokens hashed together (context of ngram_len - 1 tokens plus the candidate)
+    /// @param context_history_size number of recent contexts to remember, a repeated context is not watermarked
+    LLAMA_API struct llama_sampler * llama_sampler_init_synthid(
+                       const int64_t * keys,
+                              size_t   n_keys,
+                       const uint8_t * sampling_table,
+                              size_t   sampling_table_size,
+                             int32_t   ngram_len,
+                             int32_t   context_history_size);
+
     /// @details Mirostat 1.0 algorithm described in the paper https://arxiv.org/abs/2007.14966. Uses tokens instead of words.
     /// @param candidates A vector of `llama_token_data` containing the candidate tokens, their probabilities (p), and log-odds (logit) for the current position in the generated text.
     /// @param tau  The target cross-entropy (or surprise) value you want to achieve for the generated text. A higher value corresponds to more surprising or less predictable text, while a lower value corresponds to less surprising or more predictable text.

scripts/synthid-check.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+'''
+  Checks if a text carries the SynthID watermark produced by llama.cpp with --synthid-keys.
+  The g-values are computed with the SynthIDTextWatermarkLogitsProcessor from HF transformers,
+  the score is the mean g-value: about 50% for unwatermarked text, higher for watermarked text.
+
+  Syntax:
+    ./scripts/synthid-check.py -f output.txt --keys 654,400,836 --tokenizer <HF model id or path>
+
+  The tokenizer must be the one of the model that generated the text.
+  The text should contain only the generated part, without the prompt.
+'''
+
+import argparse
+import math
+import sys
+
+import torch
+from transformers import AutoTokenizer
+from transformers.generation.logits_process import SynthIDTextWatermarkLogitsProcessor
+
+
+def score_tokens(token_ids, keys, ngram_len=5, sampling_table_size=65536, sampling_table_seed=0, context_history_size=1024):
+    processor = SynthIDTextWatermarkLogitsProcessor(
+        ngram_len=ngram_len,
+        keys=keys,
+        sampling_table_size=sampling_table_size,
+        sampling_table_seed=sampling_table_seed,
+        context_history_size=context_history_size,
+        device=torch.device("cpu"),
+    )
+
+    ids = torch.tensor([token_ids], dtype=torch.long)
+    g_values = processor.compute_g_values(ids)                       # (1, n_ngrams, depth)
+    mask = processor.compute_context_repetition_mask(ids).float()    # (1, n_ngrams)
+
+    depth = g_values.shape[-1]
+    n_scored = int(mask.sum().item())
+    if n_scored == 0:
+        raise ValueError("text is too short to score")
+
+    mean = (g_values.float() * mask[..., None]).sum().item() / (depth * n_scored)
+
+    # unwatermarked text gives i.i.d. Bernoulli(0.5) g-values
+    z_score = (mean - 0.5) / (0.5 / math.sqrt(depth * n_scored))
+
+    return mean, z_score, n_scored
+
+
+def main():
+    parser = argparse.ArgumentParser(description="check a text for the SynthID watermark")
+    parser.add_argument("-f", "--file", required=True, help="text file to check")
+    parser.add_argument("--keys", required=True, help="comma-separated watermarking keys, same as --synthid-keys")
+    parser.add_argument("--tokenizer", required=True, help="HF model id or local path of the tokenizer")
+    parser.add_argument("--ngram-len", type=int, default=5)
+    args = parser.parse_args()
+
+    keys = [int(k) for k in args.keys.split(",") if k.strip()]
+    if not keys:
+        sys.exit("error: --keys requires at least one key")
+
+    with open(args.file, encoding="utf-8") as f:
+        text = f.read()
+
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
+    token_ids = tokenizer.encode(text, add_special_tokens=False)
+
+    if len(token_ids) < args.ngram_len:
+        sys.exit(f"error: text has {len(token_ids)} tokens, need at least {args.ngram_len}")
+
+    mean, z_score, n_scored = score_tokens(token_ids, keys, ngram_len=args.ngram_len)
+
+    print(f"tokens: {len(token_ids)}, scored ngrams: {n_scored}")
+    print(f"watermark score: {100.0 * mean:.1f}% (unwatermarked text scores about 50%)")
+    print(f"z-score: {z_score:.2f}")
+
+
+if __name__ == "__main__":
+    main()

src/llama-sampler.cpp

@@ -3312,6 +3312,186 @@ struct llama_sampler * llama_sampler_init_top_n_sigma(float n) {
     );
 }
 
+// synthid
+
+struct llama_sampler_synthid {
+    const std::vector<int64_t> keys;
+    const std::vector<uint8_t> sampling_table;
+    const int32_t ngram_len;
+    const int32_t context_history_size;
+
+    // last ngram_len - 1 accepted tokens
+    std::vector<llama_token> context;
+
+    ring_buffer<uint64_t> context_history;
+
+    std::vector<uint64_t> hashes;
+    std::vector<uint8_t>  g_values;
+
+    llama_sampler_synthid(
+            std::vector<int64_t> keys,
+            std::vector<uint8_t> sampling_table,
+            int32_t ngram_len,
+            int32_t context_history_size)
+        : keys                (std::move(keys))
+        , sampling_table      (std::move(sampling_table))
+        , ngram_len           (ngram_len)
+        , context_history_size(context_history_size)
+        , context             (ngram_len - 1, 0)
+        , context_history     (context_history_size) {
+    }
+
+    // linear congruential generator, same constants as the HF implementation
+    static uint64_t hash(uint64_t h, uint64_t x) {
+        return (h + x) * 6364136223846793005ULL + 1;
+    }
+
+    uint8_t g_value(uint64_t key) const {
+        const int64_t n = (int64_t) sampling_table.size();
+
+        // python-style modulo, the result is always in [0, n)
+        int64_t idx = (int64_t) key % n;
+        if (idx < 0) {
+            idx += n;
+        }
+
+        return sampling_table[idx];
+    }
+};
+
+static const char * llama_sampler_synthid_name(const struct llama_sampler * /*smpl*/) {
+    return "synthid";
+}
+
+static void llama_sampler_synthid_accept(struct llama_sampler * smpl, llama_token token) {
+    auto * ctx = (llama_sampler_synthid *) smpl->ctx;
+
+    if (ctx->context.empty()) {
+        return;
+    }
+
+    std::copy(ctx->context.begin() + 1, ctx->context.end(), ctx->context.begin());
+    ctx->context.back() = token;
+}
+
+static void llama_sampler_synthid_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) {
+    auto * ctx = (llama_sampler_synthid *) smpl->ctx;
+
+    uint64_t h_context = 1;
+    for (const llama_token t : ctx->context) {
+        h_context = llama_sampler_synthid::hash(h_context, (uint64_t) (int64_t) t);
+    }
+
+    bool is_repeated = false;
+    for (size_t i = 0; i < ctx->context_history.size(); ++i) {
+        if (ctx->context_history.rat(i) == h_context) {
+            is_repeated = true;
+            break;
+        }
+    }
+
+    ctx->context_history.push_back(h_context);
+
+    if (is_repeated) {
+        return;
+    }
+
+    llama_sampler_softmax_impl(cur_p, false);
+
+    ctx->hashes.resize(cur_p->size);
+    ctx->g_values.resize(cur_p->size);
+
+    for (size_t i = 0; i < cur_p->size; ++i) {
+        ctx->hashes[i] = llama_sampler_synthid::hash(h_context, (uint64_t) (int64_t) cur_p->data[i].id);
+    }
+
+    for (const int64_t key : ctx->keys) {
+        float g_mass = 0.0f;
+
+        for (size_t i = 0; i < cur_p->size; ++i) {
+            ctx->g_values[i] = ctx->g_value(llama_sampler_synthid::hash(ctx->hashes[i], (uint64_t) key));
+            g_mass += ctx->g_values[i] * cur_p->data[i].p;
+        }
+
+        for (size_t i = 0; i < cur_p->size; ++i) {
+            cur_p->data[i].p *= 1.0f + ctx->g_values[i] - g_mass;
+        }
+    }
+
+    for (size_t i = 0; i < cur_p->size; ++i) {
+        cur_p->data[i].logit = cur_p->data[i].p > 0.0f ? logf(cur_p->data[i].p) : -INFINITY;
+    }
+
+    cur_p->sorted = false;
+}
+
+static void llama_sampler_synthid_reset(struct llama_sampler * smpl) {
+    auto * ctx = (llama_sampler_synthid *) smpl->ctx;
+    std::fill(ctx->context.begin(), ctx->context.end(), 0);
+    ctx->context_history.clear();
+}
+
+static struct llama_sampler * llama_sampler_synthid_clone(const struct llama_sampler * smpl) {
+    const auto * ctx = (const llama_sampler_synthid *) smpl->ctx;
+    auto * result = llama_sampler_init_synthid(
+            ctx->keys.data(),
+            ctx->keys.size(),
+            ctx->sampling_table.data(),
+            ctx->sampling_table.size(),
+            ctx->ngram_len,
+            ctx->context_history_size);
+
+    // copy the state
+    {
+        auto * result_ctx = (llama_sampler_synthid *) result->ctx;
+
+        result_ctx->context         = ctx->context;
+        result_ctx->context_history = ctx->context_history;
+    }
+
+    return result;
+}
+
+static void llama_sampler_synthid_free(struct llama_sampler * smpl) {
+    delete (llama_sampler_synthid *) smpl->ctx;
+}
+
+static struct llama_sampler_i llama_sampler_synthid_i = {
+    /* .name              = */ llama_sampler_synthid_name,
+    /* .accept            = */ llama_sampler_synthid_accept,
+    /* .apply             = */ llama_sampler_synthid_apply,
+    /* .reset             = */ llama_sampler_synthid_reset,
+    /* .clone             = */ llama_sampler_synthid_clone,
+    /* .free              = */ llama_sampler_synthid_free,
+    /* .backend_init      = */ nullptr,
+    /* .backend_accept    = */ nullptr,
+    /* .backend_apply     = */ nullptr,
+    /* .backend_set_input = */ nullptr,
+    /* .backend_reset     = */ nullptr,
+    /* .copy_state        = */ nullptr,
+};
+
+struct llama_sampler * llama_sampler_init_synthid(
+        const int64_t * keys,
+               size_t   n_keys,
+        const uint8_t * sampling_table,
+               size_t   sampling_table_size,
+              int32_t   ngram_len,
+              int32_t   context_history_size) {
+    if (n_keys == 0 || sampling_table_size == 0 || ngram_len < 1 || context_history_size < 1) {
+        return llama_sampler_init_empty("?synthid");
+    }
+
+    return llama_sampler_init(
+        /* .iface = */ &llama_sampler_synthid_i,
+        /* .ctx   = */ new llama_sampler_synthid(
+            std::vector<int64_t>(keys, keys + n_keys),
+            std::vector<uint8_t>(sampling_table, sampling_table + sampling_table_size),
+            ngram_len,
+            context_history_size)
+    );
+}
+
 // DRY
 
 struct llama_sampler_dry {

tests/test-sampling.cpp

@@ -221,6 +221,46 @@ static void test_top_n_sigma(const std::vector<float> & probs, const std::vector
     tester.check();
 }
 
+static void test_synthid(
+    const std::vector<float> & probs, const std::vector<llama_token> & last_tokens, const std::vector<float> & probs_expected
+) {
+    GGML_ASSERT(probs.size() == probs_expected.size());
+
+    // values generated with SynthIDTextWatermarkLogitsProcessor from HF transformers
+    const std::vector<int64_t> keys = { 654, 400, 836 };
+    const std::vector<uint8_t> sampling_table = { 1, 0, 0, 1, 1, 1, 0, 1, 0, 0, 1, 0, 1, 1, 0, 1 };
+
+    auto * sampler = llama_sampler_init_synthid(keys.data(), keys.size(), sampling_table.data(), sampling_table.size(), 3, 4);
+
+    for (size_t i = 0; i < last_tokens.size(); i++) {
+        llama_sampler_accept(sampler, last_tokens[i]);
+    }
+
+    {
+        sampler_tester tester(probs, probs_expected);
+
+        DUMP(&tester.cur_p);
+        llama_sampler_apply(sampler, &tester.cur_p);
+        tester.apply(llama_sampler_init_dist(0));
+        DUMP(&tester.cur_p);
+
+        tester.check();
+    }
+
+    // same context again -> the watermark is not applied
+    {
+        sampler_tester tester(probs, probs);
+
+        llama_sampler_apply(sampler, &tester.cur_p);
+        tester.apply(llama_sampler_init_dist(0));
+        DUMP(&tester.cur_p);
+
+        tester.check();
+    }
+
+    llama_sampler_free(sampler);
+}
+
 static void test_sampler_queue(const size_t n_vocab, const std::string & samplers_sequence, const int top_k, const float top_p, const float min_p
 ) {
     sampler_tester tester(n_vocab);
@@ -395,6 +435,8 @@ int main(void) {
     test_top_n_sigma({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 0.00f); // top_n_sigma == 0 now represents a no-op rather than greedy decoding as of PR#13345
     test_top_n_sigma({0.1f, 0.2f, 0.3f, 0.4f}, {0.1f, 0.2f, 0.3f, 0.4f}, 3.00f);
 
+    test_synthid({0.05f, 0.1f, 0.15f, 0.2f, 0.25f, 0.25f}, {3, 5}, {0.049753f, 0.008964f, 0.029852f, 0.294519f, 0.368149f, 0.248764f});
+
     test_sampler_queue(10000, "k", 10000, 1.0f, 1.0f);
     test_sampler_queue(10000, "k",     1, 1.0f, 1.0f);
     test_sampler_queue(10000, "p", 10000, 1.0f, 1.0f);