add --synthid-keys

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

@@ -1396,7 +1396,7 @@ extern "C" {
 
     /// @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 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