add synthid sampler

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 one watermarking key per tournament layer
+    /// @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.

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